simplify wait

sourcenetwork · Jun 26, 2024 · 2fcac46 · 2fcac46
1 parent 51048b2
commit 2fcac46
Showing 1 changed file with 14 additions and 22 deletions.
diff --git a/tests/integration/p2p.go b/tests/integration/p2p.go
@@ -293,17 +293,15 @@ func configureReplicator(
 	err = sourceNode.SetReplicator(s.ctx, client.Replicator{
 		Info: targetNode.PeerInfo(),
 	})
+	if err == nil {
+		<-sub.Message()
+	}
 
 	expectedErrorRaised := AssertError(s.t, s.testCase.Description, err, cfg.ExpectedError)
 	assertExpectedErrorRaised(s.t, s.testCase.Description, cfg.ExpectedError, expectedErrorRaised)
 	if err == nil {
 		setupReplicatorWaitSync(s, 0, cfg)
 	}
-	for msg := range sub.Message() {
-		if msg.Name == event.ReplicatorCompletedName {
-			break
-		}
-	}
 }
 
 func deleteReplicator(
@@ -318,12 +316,10 @@ func deleteReplicator(
 	err = sourceNode.DeleteReplicator(s.ctx, client.Replicator{
 		Info: targetNode.PeerInfo(),
 	})
-	require.NoError(s.t, err)
-	for msg := range sub.Message() {
-		if msg.Name == event.ReplicatorCompletedName {
-			break
-		}
+	if err == nil {
+		<-sub.Message()
 	}
+	require.NoError(s.t, err)
 }
 
 func setupReplicatorWaitSync(
@@ -408,15 +404,13 @@ func subscribeToCollection(
 	require.NoError(s.t, err)
 
 	err = n.AddP2PCollections(s.ctx, schemaRoots)
+	if err == nil {
+		<-sub.Message()
+	}
+
 	expectedErrorRaised := AssertError(s.t, s.testCase.Description, err, action.ExpectedError)
 	assertExpectedErrorRaised(s.t, s.testCase.Description, action.ExpectedError, expectedErrorRaised)
 
-	for msg := range sub.Message() {
-		if msg.Name == event.P2PTopicCompletedName {
-			break
-		}
-	}
-
 	// The `n.Peer.AddP2PCollections(colIDs)` call above is calling some asynchronous functions
 	// for the pubsub subscription and those functions can take a bit of time to complete,
 	// we need to make sure this has finished before progressing.
@@ -447,15 +441,13 @@ func unsubscribeToCollection(
 	require.NoError(s.t, err)
 
 	err = n.RemoveP2PCollections(s.ctx, schemaRoots)
+	if err == nil {
+		<-sub.Message()
+	}
+
 	expectedErrorRaised := AssertError(s.t, s.testCase.Description, err, action.ExpectedError)
 	assertExpectedErrorRaised(s.t, s.testCase.Description, action.ExpectedError, expectedErrorRaised)
 
-	for msg := range sub.Message() {
-		if msg.Name == event.P2PTopicCompletedName {
-			break
-		}
-	}
-
 	// The `n.Peer.RemoveP2PCollections(colIDs)` call above is calling some asynchronous functions
 	// for the pubsub subscription and those functions can take a bit of time to complete,
 	// we need to make sure this has finished before progressing.