tests: make TestSendMessageCallbacks deterministic (#5523)

network/wsNetwork_test.go

@@ -4515,37 +4515,42 @@ func TestMergePrimarySecondaryRelayAddressListsNoDedupExp(t *testing.T) {
 func TestSendMessageCallbacks(t *testing.T) {
 	partitiontest.PartitionTest(t)
 	netA, netB, _, closeFunc := setupWebsocketNetworkAB(t, 2)
-	_ = netB
 	defer closeFunc()
 
 	var counter uint64
 	require.NotZero(t, netA.NumPeers())
-	peer := netA.peers[0]
-	// Need to create a channel so that the message doesn't get filtered out
-	netB.peers[0].makeResponseChannel(1)
+
+	// peerB is netA's representation of netB and vice versa
+	peerB := netA.peers[0]
+	peerA := netB.peers[0]
+
+	// Need to create a channel so that TS messages sent by netA don't get filtered out in the readLoop
+	peerA.makeResponseChannel(1)
+
+	// The for loop simulates netA receiving 100 UE block requests from netB
+	// and goes through the actual response code path to generate and send TS responses to netB
 	for i := 0; i < 100; i++ {
 		randInt := crypto.RandUint64()%(128) + 1
 		atomic.AddUint64(&counter, randInt)
 		topic := MakeTopic("val", []byte("blah"))
 		callback := func() {
 			atomic.AddUint64(&counter, ^uint64(randInt-1))
 		}
-		msg := IncomingMessage{Sender: netA.peers[0], Tag: protocol.UniEnsBlockReqTag}
-		peer.Respond(context.Background(), msg, OutgoingMessage{OnRelease: callback, Topics: Topics{topic}})
+		msg := IncomingMessage{Sender: peerB, Tag: protocol.UniEnsBlockReqTag}
+		peerB.Respond(context.Background(), msg, OutgoingMessage{OnRelease: callback, Topics: Topics{topic}})
 	}
-	// force it to disconnect by removing the only response channel -- this is breach of protocol.
-	netB.peers[0].getAndRemoveResponseChannel(1)
-	// confirm that we still have messages in the send buffer
-	require.NotZero(t, len(peer.sendBufferBulk))
-
-	// confirm that eventually the messages get drained during the cleanup
+	// Confirm that netB's representation netA peerB has received some requests and decremented the counter
+	// of outstanding TS requests below 0. This will be true because we never made any UE block requests, we only
+	// simulated them by manually creating a IncomingMessage with the UE tag in the loop above
 	require.Eventually(t,
-		func() bool { return atomic.LoadUint64(&counter) == uint64(0) },
+		func() bool { return atomic.LoadInt64(&peerA.outstandingTopicRequests) < 0 },
 		500*time.Millisecond,
 		25*time.Millisecond,
 	)
+
+	// confirm that the test counter decrements down to zero correctly through callbacks
 	require.Eventually(t,
-		func() bool { return netA.NumPeers() == 0 },
+		func() bool { return atomic.LoadUint64(&counter) == uint64(0) },
 		500*time.Millisecond,
 		25*time.Millisecond,
 	)