Refactor testcase to actively wait instead of sleep

Signed-off-by: Lukas Pukenis <[email protected]>
NordSecurity · Dec 17, 2024 · aee6ac6 · aee6ac6
1 parent dbaaf60
commit aee6ac6
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 13 deletions.
diff --git a/nat-lab/tests/telio.py b/nat-lab/tests/telio.py
@@ -114,6 +114,7 @@ async def notify_peer_state(
         paths: List[PathType],
         is_exit: bool = False,
         is_vpn: bool = False,
+        link_state: Optional[LinkState] = None,
     ) -> None:
         while True:
             peer = self.get_peer_info(public_key)
@@ -123,6 +124,7 @@ async def notify_peer_state(
                 and peer.state in states
                 and is_exit == peer.is_exit
                 and is_vpn == peer.is_vpn
+                and (link_state is None or peer.link_state == link_state)
             ):
                 return
             await asyncio.sleep(0.1)
@@ -251,9 +253,12 @@ async def wait_for_state_peer(
         is_exit: bool = False,
         is_vpn: bool = False,
         timeout: Optional[float] = None,
+        link_state: Optional[LinkState] = None,
     ) -> None:
         await asyncio.wait_for(
-            self._runtime.notify_peer_state(public_key, state, paths, is_exit, is_vpn),
+            self._runtime.notify_peer_state(
+                public_key, state, paths, is_exit, is_vpn, link_state
+            ),
             timeout,
         )
 
@@ -522,6 +527,7 @@ async def wait_for_state_peer(
         is_exit: bool = False,
         is_vpn: bool = False,
         timeout: Optional[float] = None,
+        link_state: Optional[LinkState] = None,
     ) -> None:
         await self.get_events().wait_for_state_peer(
             public_key,
@@ -530,6 +536,7 @@ async def wait_for_state_peer(
             is_exit,
             is_vpn,
             timeout,
+            link_state,
         )
 
     async def wait_for_event_peer(

diff --git a/nat-lab/tests/test_batching.py b/nat-lab/tests/test_batching.py
@@ -279,8 +279,9 @@ def features():
         )
     ],
 )
-@pytest.mark.timeout(60)
-async def test_proxying_peer_keepalive(setup_params: List[SetupParameters]) -> None:
+async def test_proxying_peer_batched_keepalive(
+    setup_params: List[SetupParameters],
+) -> None:
     # Since batching keepalives are performed on application level instead of Wireguard
     # backend we need to ensure that proxying peers are receiving the keepalives. To test
     # for that we can enable link detection that guarantees quick detection if there's no corresponding
@@ -306,13 +307,10 @@ async def test_proxying_peer_keepalive(setup_params: List[SetupParameters]) -> N
 
         _, beta_node = env.nodes
 
-        # 20 seconds should be enough since:
-        # * 10 seconds is for WireGuard's Passive-Keepalive
-        # * RTT is configured to be 2 seconds
-        # * no pings to validate
-        # * PersistentKeepalive configured to 5seconds
-        # Given all these it should be worst case latency until detection: 10+2+5=17
-        await asyncio.sleep(20)
-
-        link_events = alpha.get_link_state_events(beta_node.public_key)
-        assert link_events[-1] == LinkState.DOWN
+        await alpha.wait_for_state_peer(
+            beta_node.public_key,
+            [NodeState.CONNECTED],
+            [PathType.RELAY],
+            timeout=30,
+            link_state=LinkState.DOWN,
+        )