diff --git a/nat-lab/tests/test_direct_connection.py b/nat-lab/tests/test_direct_connection.py
index 59b478907..11520551c 100644
--- 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)
diff --git a/nat-lab/tests/uniffi/libtelio_proxy.py b/nat-lab/tests/uniffi/libtelio_proxy.py
index dfccad918..dcf122aff 100644
--- a/nat-lab/tests/uniffi/libtelio_proxy.py
+++ b/nat-lab/tests/uniffi/libtelio_proxy.py
@@ -56,7 +56,7 @@ def shutdown(self, container_or_vm_name: Optional[str] = None):
                 "Unknown" if container_or_vm_name is None else container_or_vm_name,
             )
 
-        except Pyro5.errors.ConnectionClosedError as e:
+        except (Pyro5.errors.ConnectionClosedError, Pyro5.errors.ConnectionRefusedError) as e:
             # Shutting down the server via client request is naturally racy,
             # as sending of response is racing against process shutdown (and
             # thus server-side socket being closed).