Remove flakiness from mina_net2 tests

src/lib/mina_net2/tests/all_ipc.ml

@@ -197,15 +197,6 @@ let%test_module "all-ipc test" =
                 raise UnexpectedState)
         |> or_timeout ~msg:"Alice: wait for Carol to connect"
       in
-      (* List peers of Alice *)
-      let%bind peers = peers a in
-      assert (List.length peers = 3) ;
-      assert (
-        List.fold [ ad.y_peerid; ad.b_peerid; ad.c_peerid ] ~init:true
-          ~f:(fun b_acc pid ->
-            b_acc
-            && List.fold peers ~init:false ~f:(fun acc p ->
-                   acc || String.equal p.peer_id pid)) ) ;
       (* Subscribe to topic "c" *)
       let topic_c_received_1 = ref false in
       let topic_c_received_2 = ref false in
@@ -327,6 +318,16 @@ let%test_module "all-ipc test" =
          (waiting for Alice to disconnect) *)
       let%bind () = reset_stream a stream2 >>| Or_error.ok_exn in
 
+      (* List peers of Alice *)
+      let%bind peers = peers a in
+      assert (List.length peers >= 2) ;
+      assert (
+        List.fold [ ad.y_peerid; ad.b_peerid; ad.c_peerid ] ~init:true
+          ~f:(fun b_acc pid ->
+            b_acc
+            && List.fold peers ~init:false ~f:(fun acc p ->
+                   acc || String.equal p.peer_id pid)) ) ;
+
       (* Ban Carol in Alice's gating config *)
       let%bind _ =
         set_connection_gating_config a (mk_banning_gating_config ad.c_peerid)
@@ -577,7 +578,7 @@ let%test_module "all-ipc test" =
       in
       return (node, peerid, addr, shutdown)
 
-    let test_def =
+    let test_def () =
       let open Deferred.Let_syntax in
       let on_connected (_, w) s = don't_wait_for (Pipe.write w (true, s)) in
       let on_disconnected (_, w) s = don't_wait_for (Pipe.write w (false, s)) in
@@ -640,5 +641,5 @@ let%test_module "all-ipc test" =
     let%test_unit "ipc test" =
       (* ignore test_def *)
       let () = Core.Backtrace.elide := false in
-      Async.Thread_safe.block_on_async_exn (fun () -> test_def)
+      Async.Thread_safe.block_on_async_exn test_def
   end )

src/lib/mina_net2/tests/tests.ml

@@ -86,6 +86,7 @@ let%test_module "coda network tests" =
           ~known_private_ip_nets:[]
         >>| Or_error.ok_exn
       in
+      let%bind () = after (Time.Span.of_sec 10.) in
       let%bind b_advert = begin_advertising b in
       Or_error.ok_exn b_advert ;
       let%bind c_advert = begin_advertising c in
@@ -105,7 +106,7 @@ let%test_module "coda network tests" =
     (* TODO fails occasionally, uncomment after debugging it *)
     let%test_unit "b_stream_c" =
       let () = Core.Backtrace.elide := false in
-      let test_def =
+      let test_def () =
         let open Deferred.Let_syntax in
         let%bind b, c, shutdown = setup_two_nodes "test_stream" in
         let%bind b_peerid = me b >>| Keypair.to_peer_id in
@@ -160,11 +161,11 @@ let%test_module "coda network tests" =
         let%bind _msgs = Pipe.read_all r in
         shutdown ()
       in
-      Async.Thread_safe.block_on_async_exn (fun () -> test_def)
+      Async.Thread_safe.block_on_async_exn test_def
 
     let%test_unit "stream" =
       let () = Core.Backtrace.elide := false in
-      let test_def =
+      let test_def () =
         let open Deferred.Let_syntax in
         let%bind b, c, shutdown = setup_two_nodes "test_stream" in
         let%bind b_peerid = me b >>| Keypair.to_peer_id in
@@ -193,8 +194,7 @@ let%test_module "coda network tests" =
         assert (String.equal msg testmsg) ;
         let%bind () = Ivar.read handler_finished in
         let%bind () = close_protocol b ~protocol:"echo" in
-        let%map () = shutdown () in
-        ()
+        shutdown ()
       in
-      Async.Thread_safe.block_on_async_exn (fun () -> test_def)
+      Async.Thread_safe.block_on_async_exn test_def
   end )