From 70daa7812bac8df19daebd26c2a89594c4fc56b0 Mon Sep 17 00:00:00 2001
From: Guilherme Ferreira <guilherme.ferreira@protofire.io>
Date: Fri, 19 Apr 2024 10:27:03 -0300
Subject: [PATCH] Add more keys to GetSize test

---
 p2p/peerManager/peerdb/datastore_test.go | 87 +++++++++++++++++++++---
 1 file changed, 78 insertions(+), 9 deletions(-)

diff --git a/p2p/peerManager/peerdb/datastore_test.go b/p2p/peerManager/peerdb/datastore_test.go
index 14bda979a3..b01f6bcf3b 100644
--- a/p2p/peerManager/peerdb/datastore_test.go
+++ b/p2p/peerManager/peerdb/datastore_test.go
@@ -3,6 +3,7 @@ package peerdb
 import (
 	"context"
 	"encoding/json"
+	"sync"
 	"testing"
 
 	datastore "github.com/ipfs/go-datastore"
@@ -113,23 +114,91 @@ func TestQuery(t *testing.T) {
 
 func TestGetSize(t *testing.T) {
 	ps, teardown := setupDB(t)
-	t.Cleanup(teardown)
+	defer teardown()
 
-	peer := createPeers(t, 1)[0]
+	peers := createPeers(t, 5)
 
-	key := datastore.NewKey(peer.AddrInfo.ID.String())
+	keys := make([]datastore.Key, len(peers))
+
+	cases := []struct {
+		i       int
+		Entropy uint64
+		PubKey  []byte
+	}{
+		{
+			i:       0,
+			Entropy: uint64(12345),
+			PubKey:  []byte(""),
+		},
+		{
+			i:       1,
+			Entropy: uint64(1234567890),
+			PubKey:  []byte("pub"),
+		},
+		{
+			i:       2,
+			Entropy: uint64(1),
+			PubKey:  []byte("pubkey"),
+		},
+		{
+			i:       3,
+			Entropy: uint64(0),
+			PubKey:  []byte("pubkey1234567890"),
+		},
+		{
+			i:       4,
+			Entropy: uint64(12345678901234),
+			PubKey:  []byte("pubkey12345678901234567890"),
+		},
+	}
+
+	var wg sync.WaitGroup
+
+	// Add value to keys and test first time
+	for i, peer := range peers {
+		wg.Add(1)
+		go func(peer *PeerInfo, i int) {
+			defer wg.Done()
+			keys[i] = datastore.NewKey(peer.AddrInfo.ID.String())
+			value, err := json.Marshal(peer)
+			require.NoError(t, err)
+
+			err = ps.Put(context.Background(), keys[i], value)
+			require.NoError(t, err)
+			size, err := ps.GetSize(context.Background(), keys[i])
+			require.NoError(t, err)
+			require.Equal(t, len(value), size)
+		}(peer, i)
+	}
+	wg.Wait()
+
+	// Update keys is parallel and check if size is updated
+	for i, peer := range peers {
+		wg.Add(1)
+		go func(peer *PeerInfo, i int) {
+			defer wg.Done()
+			peer.Entropy = cases[i].Entropy
+			testSize(t, ps, peer, keys[i])
+
+			peer.PubKey = cases[i].PubKey
+			testSize(t, ps, peer, keys[i])
+		}(peer, i)
+	}
+	wg.Wait()
+
+	//Test with non existent key
+	size, err := ps.GetSize(context.Background(), datastore.NewKey("non-existent-key"))
+	require.Error(t, err)
+	require.Equal(t, 0, size)
+}
+
+func testSize(t *testing.T, ps *PeerDB, peer *PeerInfo, key datastore.Key) {
 	value, err := json.Marshal(peer)
 	require.NoError(t, err)
 
 	err = ps.Put(context.Background(), key, value)
 	require.NoError(t, err)
-
 	size, err := ps.GetSize(context.Background(), key)
 	require.NoError(t, err)
 	require.Equal(t, len(value), size)
-
-	//Test with non existent key
-	size, err = ps.GetSize(context.Background(), datastore.NewKey("non-existent-key"))
-	require.Error(t, err)
-	require.Equal(t, 0, size)
 }