From 778cf98b89746a53ccf61c4ee80af022acb70912 Mon Sep 17 00:00:00 2001
From: Igor Sirotin <igor.sirotin.1012@gmail.com>
Date: Sun, 14 Jul 2024 23:42:41 +0100
Subject: [PATCH] test: peer disconnect

---
 internal/transport/waku_test.go | 39 ++++++++++++++++++++-------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/internal/transport/waku_test.go b/internal/transport/waku_test.go
index 108be1f..54bff16 100644
--- a/internal/transport/waku_test.go
+++ b/internal/transport/waku_test.go
@@ -102,21 +102,30 @@ func (s *WakuSuite) TestWatchConnectionStatus() {
 		close(finished)
 	}()
 
-	sent := node.PeerConnection{
-		PeerID:    peer.ID(gofakeit.UUID()),
-		Connected: true,
-	}
-
-	s.node.peerConnection <- sent
-
-	select {
-	case received := <-sub:
-		s.Require().True(received.IsOnline)
-		s.Require().False(received.HasHistory)
-		s.Require().Equal(1, received.PeersCount)
-		s.Require().True(reflect.DeepEqual(received, s.node.ConnectionStatus()))
-	case <-time.After(500 * time.Millisecond):
-		s.Require().Fail("timeout waiting for connection status")
+	peerID := peer.ID(gofakeit.UUID())
+
+	for _, connected := range []bool{true, false} {
+		expectedCount := 0
+		if connected {
+			expectedCount = 1
+		}
+
+		sent := node.PeerConnection{
+			PeerID:    peerID,
+			Connected: connected,
+		}
+
+		s.node.peerConnection <- sent
+
+		select {
+		case received := <-sub:
+			s.Require().Equal(connected, received.IsOnline)
+			s.Require().False(received.HasHistory)
+			s.Require().Equal(expectedCount, received.PeersCount)
+			s.Require().True(reflect.DeepEqual(received, s.node.ConnectionStatus()))
+		case <-time.After(500 * time.Millisecond):
+			s.Require().Fail("timeout waiting for connection status")
+		}
 	}
 
 	close(s.node.peerConnection)