From 38078489d658e6011e6b3369c0585e6f5358a62a Mon Sep 17 00:00:00 2001
From: Vlad <13818348+walldiss@users.noreply.github.com>
Date: Wed, 28 Aug 2024 20:42:37 +0200
Subject: [PATCH] fix TestDiscovery test

---
 share/p2p/discovery/discovery_test.go | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/share/p2p/discovery/discovery_test.go b/share/p2p/discovery/discovery_test.go
index 8214a2bbe0..8fb31de922 100644
--- a/share/p2p/discovery/discovery_test.go
+++ b/share/p2p/discovery/discovery_test.go
@@ -28,7 +28,7 @@ func TestDiscovery(t *testing.T) {
 
 	discoveryRetryTimeout = time.Millisecond * 100 // defined in discovery.go
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*20)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*40)
 	t.Cleanup(cancel)
 
 	tn := newTestnet(ctx, t)
@@ -50,21 +50,23 @@ func TestDiscovery(t *testing.T) {
 	peerA := tn.startNewDiscovery(params, host, routingDisc, fullNodesTag,
 		WithOnPeersUpdate(submit),
 	)
-
-	// start discovery advertisement services for other peers
 	params.AdvertiseInterval = time.Millisecond * 100
-	discs := make([]*Discovery, nodes)
-	for i := range discs {
+	// start discovery advertisement services for other peers
+	discs := make(map[string]*Discovery, nodes)
+	for range nodes {
 		host, routingDisc := tn.peer()
-		disc, err := NewDiscovery(params, host, routingDisc, fullNodesTag)
-		require.NoError(t, err)
-		go disc.Advertise(tn.ctx)
-		discs[i] = tn.startNewDiscovery(params, host, routingDisc, fullNodesTag)
+		disc := tn.startNewDiscovery(params, host, routingDisc, fullNodesTag, WithAdvertise())
+		discs[disc.host.ID().String()] = disc
+	}
 
+	for range nodes {
 		select {
 		case res := <-updateCh:
-			require.Equal(t, discs[i].host.ID(), res.peerID)
 			require.True(t, res.isAdded)
+			if _, ok := discs[res.peerID.String()]; !ok {
+				t.Fatal("discovered unknown peer")
+			}
+			delete(discs, res.peerID.String())
 		case <-ctx.Done():
 			t.Fatal("did not discover peer in time")
 		}
@@ -87,7 +89,7 @@ func TestDiscovery(t *testing.T) {
 		}
 	}
 
-	assert.EqualValues(t, 0, peerA.set.Size())
+	assert.Equal(t, uint(10), peerA.set.Size())
 }
 
 func TestDiscoveryTagged(t *testing.T) {
@@ -190,8 +192,6 @@ func (t *testnet) peer() (host.Host, discovery.Discovery) {
 	dht, err := dht.New(t.ctx, hst,
 		dht.Mode(dht.ModeServer),
 		dht.ProtocolPrefix("/test"),
-		// needed to reduce connections to peers on DHT level
-		dht.BucketSize(1),
 	)
 	require.NoError(t.T, err)