From 6eada787d951770616bc31ea749c7dcebf935a08 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 24 Sep 2021 16:16:16 +0100
Subject: [PATCH 1/2] fix flaky TestFastDisconnect identify test

---
 p2p/protocol/identify/id_glass_test.go | 32 +++++++++++---------------
 1 file changed, 14 insertions(+), 18 deletions(-)

diff --git a/p2p/protocol/identify/id_glass_test.go b/p2p/protocol/identify/id_glass_test.go
index 8d7bf5eb..1924da09 100644
--- a/p2p/protocol/identify/id_glass_test.go
+++ b/p2p/protocol/identify/id_glass_test.go
@@ -5,12 +5,13 @@ import (
 	"testing"
 	"time"
 
+	blhost "github.com/libp2p/go-libp2p-blankhost"
 	"github.com/libp2p/go-libp2p-core/network"
 	"github.com/libp2p/go-libp2p-core/peer"
-	"github.com/stretchr/testify/require"
-
-	blhost "github.com/libp2p/go-libp2p-blankhost"
 	swarmt "github.com/libp2p/go-libp2p-swarm/testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestFastDisconnect(t *testing.T) {
@@ -27,7 +28,7 @@ func TestFastDisconnect(t *testing.T) {
 
 	sync := make(chan struct{})
 	target.SetStreamHandler(ID, func(s network.Stream) {
-		// Wait till the stream is setup on both sides.
+		// Wait till the stream is set up on both sides.
 		select {
 		case <-sync:
 		case <-ctx.Done():
@@ -36,10 +37,11 @@ func TestFastDisconnect(t *testing.T) {
 
 		// Kill the connection, and make sure we're completely disconnected.
 		s.Conn().Close()
-		for target.Network().Connectedness(s.Conn().RemotePeer()) == network.Connected {
-			// let something else run
-			time.Sleep(time.Millisecond)
-		}
+		assert.Eventually(t,
+			func() bool { return target.Network().Connectedness(s.Conn().RemotePeer()) != network.Connected },
+			2*time.Second,
+			time.Millisecond,
+		)
 		// Now try to handle the response.
 		// This should not block indefinitely, or panic, or anything like that.
 		//
@@ -57,14 +59,10 @@ func TestFastDisconnect(t *testing.T) {
 	source := blhost.NewBlankHost(swarmt.GenSwarm(t))
 	defer source.Close()
 
-	err = source.Connect(ctx, peer.AddrInfo{ID: target.ID(), Addrs: target.Addrs()})
-	if err != nil {
-		t.Fatal(err)
-	}
+	// only connect to the first address, to make sure we only end up with one connection
+	require.NoError(t, source.Connect(ctx, peer.AddrInfo{ID: target.ID(), Addrs: target.Addrs()[:1]}))
 	s, err := source.NewStream(ctx, target.ID(), ID)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	select {
 	case sync <- struct{}{}:
 	case <-ctx.Done():
@@ -77,7 +75,5 @@ func TestFastDisconnect(t *testing.T) {
 		t.Fatal(ctx.Err())
 	}
 	// double-check to make sure we didn't actually timeout somewhere.
-	if ctx.Err() != nil {
-		t.Fatal(ctx.Err())
-	}
+	require.NoError(t, ctx.Err())
 }

From 7fa07a4aacd5f8148ab6b73c89346689ad9b0d3f Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 24 Sep 2021 18:49:40 +0100
Subject: [PATCH 2/2] close all connections in the TestFastDisconnect test

---
 p2p/protocol/identify/id_glass_test.go | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/p2p/protocol/identify/id_glass_test.go b/p2p/protocol/identify/id_glass_test.go
index 1924da09..273391fe 100644
--- a/p2p/protocol/identify/id_glass_test.go
+++ b/p2p/protocol/identify/id_glass_test.go
@@ -36,9 +36,13 @@ func TestFastDisconnect(t *testing.T) {
 		}
 
 		// Kill the connection, and make sure we're completely disconnected.
-		s.Conn().Close()
 		assert.Eventually(t,
-			func() bool { return target.Network().Connectedness(s.Conn().RemotePeer()) != network.Connected },
+			func() bool {
+				for _, conn := range target.Network().ConnsToPeer(s.Conn().RemotePeer()) {
+					conn.Close()
+				}
+				return target.Network().Connectedness(s.Conn().RemotePeer()) != network.Connected
+			},
 			2*time.Second,
 			time.Millisecond,
 		)
@@ -60,7 +64,7 @@ func TestFastDisconnect(t *testing.T) {
 	defer source.Close()
 
 	// only connect to the first address, to make sure we only end up with one connection
-	require.NoError(t, source.Connect(ctx, peer.AddrInfo{ID: target.ID(), Addrs: target.Addrs()[:1]}))
+	require.NoError(t, source.Connect(ctx, peer.AddrInfo{ID: target.ID(), Addrs: target.Addrs()}))
 	s, err := source.NewStream(ctx, target.ID(), ID)
 	require.NoError(t, err)
 	select {