From e49dfd7824c35c99ac92cfaa94ec020d424f02d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CGheisMohammadi=E2=80=9D?=
 <36589218+GheisMohammadi@users.noreply.github.com>
Date: Thu, 12 Jan 2023 08:41:46 +0800
Subject: [PATCH] fix p2p discovery test issue

---
 p2p/discovery/discovery_test.go                  | 10 +++++++++-
 p2p/discovery/option_test.go                     |  6 +++---
 p2p/stream/common/streammanager/streammanager.go |  2 +-
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/p2p/discovery/discovery_test.go b/p2p/discovery/discovery_test.go
index 6ea902375..f2c2f5b99 100644
--- a/p2p/discovery/discovery_test.go
+++ b/p2p/discovery/discovery_test.go
@@ -3,9 +3,11 @@ package discovery
 // TODO: test this module
 
 import (
+	"context"
 	"testing"
 
 	"github.com/libp2p/go-libp2p"
+	dht "github.com/libp2p/go-libp2p-kad-dht"
 )
 
 func TestNewDHTDiscovery(t *testing.T) {
@@ -13,7 +15,13 @@ func TestNewDHTDiscovery(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	_, err = NewDHTDiscovery(host, DHTConfig{})
+	ctx, cancel := context.WithCancel(context.Background())
+	var idht *dht.IpfsDHT
+	idht, err = dht.New(ctx, host)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = NewDHTDiscovery(ctx, cancel, host, idht, DHTConfig{})
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/p2p/discovery/option_test.go b/p2p/discovery/option_test.go
index 747d7ca95..e782151f5 100644
--- a/p2p/discovery/option_test.go
+++ b/p2p/discovery/option_test.go
@@ -40,14 +40,14 @@ func TestDHTOption_getLibp2pRawOptions(t *testing.T) {
 			opt: DHTConfig{
 				BootNodes: testAddrStr,
 			},
-			expLen: 1,
+			expLen: 2,
 		},
 		{
 			opt: DHTConfig{
 				BootNodes:     testAddrStr,
 				DataStoreFile: &validPath,
 			},
-			expLen: 2,
+			expLen: 3,
 		},
 		{
 			opt: DHTConfig{
@@ -58,7 +58,7 @@ func TestDHTOption_getLibp2pRawOptions(t *testing.T) {
 		},
 	}
 	for i, test := range tests {
-		opts, err := test.opt.getLibp2pRawOptions()
+		opts, err := test.opt.GetLibp2pRawOptions()
 		if assErr := assertError(err, test.expErr); assErr != nil {
 			t.Errorf("Test %v: %v", i, assErr)
 		}
diff --git a/p2p/stream/common/streammanager/streammanager.go b/p2p/stream/common/streammanager/streammanager.go
index 064cb0c2d..98e20cbaf 100644
--- a/p2p/stream/common/streammanager/streammanager.go
+++ b/p2p/stream/common/streammanager/streammanager.go
@@ -323,7 +323,7 @@ func (sm *streamManager) discoverAndSetupStream(discCtx context.Context) (int, e
 			// The ctx here is using the module context instead of discover context
 			err := sm.setupStreamWithPeer(sm.ctx, pid)
 			if err != nil {
-				sm.coolDownCache.Add(peer.ID)
+				sm.coolDownCache.Add(pid)
 				sm.logger.Warn().Err(err).Str("peerID", string(pid)).Msg("failed to setup stream with peer")
 				return
 			}