p2p/simulations: TestSnapshot fail if Load function returns early

p2p/simulations/network_test.go

@@ -175,24 +175,23 @@ OUTER:
 	// subscribe to peer events
 	// every node up and conn up event will generate one additional control event
 	// therefore multiply the count by two
-	evC = make(chan *Event, (len(snap.Conns)*2)+(len(snap.Nodes)*2))
+	evC = make(chan *Event)
 	sub = network.Events().Subscribe(evC)
 	defer sub.Unsubscribe()
 
-	// load the snapshot
+	// Channel that is signalling when the Load function returns
-	// spawn separate thread to avoid deadlock in the event listeners
+	// to fail the test in the event for loop if that happens
-	err = network.Load(snap)
+	// before all connection events are counted.
-	if err != nil {
+	loadDoneC := make(chan struct{})
-		t.Fatal(err)
-	}
 
+	go func() {
 		// collect connection events up to expected number
 		ctx, cancel = context.WithTimeout(context.TODO(), time.Second*3)
 		defer cancel()
 
 		connEventCount = nodeCount
 
-OUTER_TWO:
+	OUTER_TWO:
 		for {
 			select {
 			case <-ctx.Done():
@@ -213,6 +212,10 @@ OUTER_TWO:
 						break OUTER_TWO
 					}
 				}
+			case <-loadDoneC:
+				// if load function returns before all expected events are caught.
+				// fail the test
+				t.Fatal("load function returned before all connections are established")
 			}
 		}
 
@@ -234,6 +237,16 @@ OUTER_TWO:
 				t.Fatalf("network missing conn %v -> %v", snapConn.One, snapConn.Other)
 			}
 		}
+	}()
+
+	// load the snapshot
+	// spawn separate thread to avoid deadlock in the event listeners
+	err = network.Load(snap)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// signal the event for event loop that Load function has returned
+	close(loadDoneC)
 
 	// verify that network didn't generate any other additional connection events after the ones we have collected within a reasonable period of time
 	ctx, cancel = context.WithTimeout(context.TODO(), time.Second)
@@ -241,7 +254,7 @@ OUTER_TWO:
 	select {
 	case <-ctx.Done():
 	case ev := <-evC:
-		if ev.Type == EventTypeConn {
+		if ev.Type == EventTypeConn && !ev.Control {
 			t.Fatalf("Superfluous conn found %v -> %v", ev.Conn.One, ev.Conn.Other)
 		}
 	}