autonat: fix closing of listeners in dialPolicy tests (#2226)

libp2p · Apr 5, 2023 · 88132e9 · 88132e9
1 parent 8c466ab
commit 88132e9
Showing 1 changed file with 16 additions and 9 deletions.
diff --git a/p2p/host/autonat/dialpolicy_test.go b/p2p/host/autonat/dialpolicy_test.go
@@ -23,31 +23,33 @@ func makeMA(a string) multiaddr.Multiaddr {
 }
 
 type mockT struct {
-	ctx  context.Context
-	addr multiaddr.Multiaddr
+	ctx    context.Context
+	cancel context.CancelFunc
+	addr   multiaddr.Multiaddr
 }
 
 func (m *mockT) Dial(ctx context.Context, a multiaddr.Multiaddr, p peer.ID) (transport.CapableConn, error) {
 	return nil, nil
 }
 func (m *mockT) CanDial(_ multiaddr.Multiaddr) bool { return true }
 func (m *mockT) Listen(a multiaddr.Multiaddr) (transport.Listener, error) {
-	return &mockL{m.ctx, m.addr}, nil
+	return &mockL{m.ctx, m.cancel, m.addr}, nil
 }
 func (m *mockT) Protocols() []int { return []int{multiaddr.P_IP4} }
 func (m *mockT) Proxy() bool      { return false }
 func (m *mockT) String() string   { return "mock-tcp-ipv4" }
 
 type mockL struct {
-	ctx  context.Context
-	addr multiaddr.Multiaddr
+	ctx    context.Context
+	cancel context.CancelFunc
+	addr   multiaddr.Multiaddr
 }
 
 func (l *mockL) Accept() (transport.CapableConn, error) {
 	<-l.ctx.Done()
 	return nil, errors.New("expected in mocked test")
 }
-func (l *mockL) Close() error                   { return nil }
+func (l *mockL) Close() error                   { l.cancel(); return nil }
 func (l *mockL) Addr() net.Addr                 { return nil }
 func (l *mockL) Multiaddr() multiaddr.Multiaddr { return l.addr }
 
@@ -56,6 +58,7 @@ func TestSkipDial(t *testing.T) {
 	defer cancel()
 
 	s := swarmt.GenSwarm(t)
+	defer s.Close()
 	d := dialPolicy{host: blankhost.NewBlankHost(s)}
 	if d.skipDial(makeMA("/ip4/8.8.8.8")) != false {
 		t.Fatal("failed dialing a valid public addr")
@@ -69,7 +72,7 @@ func TestSkipDial(t *testing.T) {
 		t.Fatal("didn't skip dialing an internal addr")
 	}
 
-	s.AddTransport(&mockT{ctx, makeMA("/ip4/8.8.8.8")})
+	s.AddTransport(&mockT{ctx, cancel, makeMA("/ip4/8.8.8.8")})
 	err := s.AddListenAddr(makeMA("/ip4/8.8.8.8"))
 	if err != nil {
 		t.Fatal(err)
@@ -84,6 +87,8 @@ func TestSkipPeer(t *testing.T) {
 	defer cancel()
 
 	s := swarmt.GenSwarm(t)
+	defer s.Close()
+
 	d := dialPolicy{host: blankhost.NewBlankHost(s)}
 	if d.skipPeer([]multiaddr.Multiaddr{makeMA("/ip4/8.8.8.8")}) != false {
 		t.Fatal("failed dialing a valid public addr")
@@ -95,7 +100,7 @@ func TestSkipPeer(t *testing.T) {
 		t.Fatal("succeeded with no public addr")
 	}
 
-	s.AddTransport(&mockT{ctx, makeMA("/ip4/8.8.8.8")})
+	s.AddTransport(&mockT{ctx, cancel, makeMA("/ip4/8.8.8.8")})
 	err := s.AddListenAddr(makeMA("/ip4/8.8.8.8"))
 	if err != nil {
 		t.Fatal(err)
@@ -117,8 +122,10 @@ func TestSkipLocalPeer(t *testing.T) {
 	defer cancel()
 
 	s := swarmt.GenSwarm(t)
+	defer s.Close()
+
 	d := dialPolicy{host: blankhost.NewBlankHost(s)}
-	s.AddTransport(&mockT{ctx, makeMA("/ip4/192.168.0.1")})
+	s.AddTransport(&mockT{ctx, cancel, makeMA("/ip4/192.168.0.1")})
 	err := s.AddListenAddr(makeMA("/ip4/192.168.0.1"))
 	if err != nil {
 		t.Fatal(err)