Refactor test

NordSecurity · Jan 10, 2025 · cef1da4 · cef1da4
1 parent 70bc364
commit cef1da4
Showing 1 changed file with 51 additions and 12 deletions.
diff --git a/nat-lab/tests/test_start_with_tun.py b/nat-lab/tests/test_start_with_tun.py
@@ -1,8 +1,12 @@
 from contextlib import AsyncExitStack
-from helpers import SetupParameters, ping_between_all_nodes, setup_mesh_nodes
+from helpers import SetupParameters, ping_between_all_nodes, setup_api, setup_connections, setup_mesh_nodes
+from mesh_api import API
+from telio import Client
 from typing import List
 from utils.bindings import default_features, TelioAdapterType
-from utils.connection_util import ConnectionTag
+from utils.connection_util import ConnectionTag, generate_connection_tracker_config, new_connection_with_conn_tracker
+from utils.ping import ping
+from utils.router import IPStack
 
 
 def _generate_setup_parameters(
@@ -26,13 +30,48 @@ async def test_start_with_tun() -> None:
     ])
 
     async with AsyncExitStack() as exit_stack:
-        env = await setup_mesh_nodes(exit_stack, setup_params)
-        alpha_client, _ = env.clients
-        alpha, _ = env.nodes
-
-        await alpha_client.stop_device()
-        tun = await alpha_client.create_tun("tun11")
-        await alpha_client.start_with_tun(tun, "tun11")
-        await alpha_client.simple_start()
-        await alpha_client.set_meshnet_config(env.api.get_meshnet_config(alpha.id))
-        await ping_between_all_nodes(env)
+        alpha_connection_tag = ConnectionTag.DOCKER_CONE_CLIENT_1
+        beta_connection_tag = ConnectionTag.DOCKER_CONE_CLIENT_2
+
+        api = API()
+
+        (alpha, beta) = api.default_config_two_nodes(
+            alpha_ip_stack=IPStack.IPv4v6,
+            beta_ip_stack=IPStack.IPv4v6,
+            allow_peer_traffic_routing=True,
+        )
+        (connection_alpha, alpha_conn_tracker) = await exit_stack.enter_async_context(
+            new_connection_with_conn_tracker(
+                alpha_connection_tag,
+                generate_connection_tracker_config(
+                    alpha_connection_tag,
+                    derp_1_limits=(1, 1),
+                    ping6_limits=(None, None),
+                ),
+            )
+        )
+        (connection_beta, beta_conn_tracker) = await exit_stack.enter_async_context(
+            new_connection_with_conn_tracker(
+                beta_connection_tag,
+                generate_connection_tracker_config(
+                    beta_connection_tag,
+                    derp_1_limits=(1, 1),
+                    ping6_limits=(None, None),
+                ),
+            )
+        )
+
+        client_alpha = await exit_stack.enter_async_context(
+            Client(connection_alpha, alpha, TelioAdapterType.BORING_TUN).run()
+        )
+
+        await exit_stack.enter_async_context(
+            Client(connection_beta, beta, TelioAdapterType.BORING_TUN).run()
+        )
+
+        await client_alpha.stop_device()
+        tun = await client_alpha.create_tun("tun11")
+        await client_alpha.start_with_tun(tun, "tun11")
+        await client_alpha.simple_start()
+        await client_alpha.set_meshnet_config(api.get_meshnet_config(alpha.id))
+        await ping(connection_alpha, beta.ip_addresses[0])