Bring back stricter events assertion in test-direct-working-path test…

… in order to enable easier reproduction of the issue
NordSecurity · Nov 12, 2024 · b4f3b13 · b4f3b13
1 parent f4920a9
commit b4f3b13
Showing 1 changed file with 24 additions and 18 deletions.
diff --git a/nat-lab/tests/test_direct_connection.py b/nat-lab/tests/test_direct_connection.py
@@ -177,33 +177,39 @@ async def test_direct_working_paths(
 
         print(datetime.now(), "Test direct connection")
         await _check_if_true_direct_connection(env)
+
+        print(datetime.now(), "Downgrade to relay connections")
+        # Start collection just before breaking, to not miss disconnects
+        relay_events = await asyncio.gather(*[
+            exit_stack.enter_async_context(
+                run_async_context(
+                    client.wait_for_event_peer(
+                        node.public_key, [NodeState.CONNECTED], [PathType.RELAY]
+                    )
+                )
+            )
+            for client, node in itertools.product(env.clients, env.nodes)
+            if not client.is_node(node)
+        ])
         async with _disable_direct_connection(env, reflexive_ips):
-            print(datetime.now(), "Downgrade to relay connections")
-            await asyncio.gather(*[
-                await exit_stack.enter_async_context(
+            await asyncio.gather(*relay_events)
+            await ping_between_all_nodes(env)
+
+            # Start collecting direct events, before drop of blocking
+            print(datetime.now(), "Reconnect to direct connections")
+            direct_events = await asyncio.gather(*[
+                exit_stack.enter_async_context(
                     run_async_context(
-                        client.wait_for_state_peer(
-                            node.public_key, [NodeState.CONNECTED], [PathType.RELAY]
+                        client.wait_for_event_peer(
+                            node.public_key, [NodeState.CONNECTED], [PathType.DIRECT]
                         )
                     )
                 )
                 for client, node in itertools.product(env.clients, env.nodes)
                 if not client.is_node(node)
             ])
-            await ping_between_all_nodes(env)
 
-        print(datetime.now(), "Reconnect to direct connections")
-        await asyncio.gather(*[
-            await exit_stack.enter_async_context(
-                run_async_context(
-                    client.wait_for_state_peer(
-                        node.public_key, [NodeState.CONNECTED], [PathType.DIRECT]
-                    )
-                )
-            )
-            for client, node in itertools.product(env.clients, env.nodes)
-            if not client.is_node(node)
-        ])
+        await asyncio.gather(*direct_events)
 
         print(datetime.now(), "Test direct connection again")
         await _check_if_true_direct_connection(env)