fix: preserve tunnel client test hook

Author: d42me

packages/prime/src/prime_cli/commands/tunnel.py

@@ -1,6 +1,6 @@
 import asyncio
 import signal
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import typer
 from rich.table import Table
@@ -18,6 +18,16 @@
 app = PlainTyper(help="Manage tunnels for exposing local services", no_args_is_help=True)
 console = get_console()
 
+TunnelClient: Any = None
+
+
+def _tunnel_client_class() -> Any:
+    if TunnelClient is not None:
+        return TunnelClient
+    from prime_tunnel.core.client import TunnelClient as loaded_tunnel_client
+
+    return loaded_tunnel_client
+
 
 def _format_tunnel_for_output(tunnel) -> dict:
     created_at = tunnel.created_at
@@ -157,9 +167,7 @@ def list_tunnels(
     validate_output_format(output, console)
 
     async def fetch_tunnels():
-        from prime_tunnel.core.client import TunnelClient
-
-        client = TunnelClient()
+        client = _tunnel_client_class()()
         try:
             return await client.list_tunnels_page(
                 team_id=team_id,
@@ -248,9 +256,7 @@ def tunnel_status(
     """Get status of a specific tunnel."""
 
     async def fetch_status():
-        from prime_tunnel.core.client import TunnelClient
-
-        client = TunnelClient()
+        client = _tunnel_client_class()()
         try:
             return await client.get_tunnel(tunnel_id)
         finally:
@@ -336,9 +342,7 @@ def stop_tunnel(
     if all:
 
         async def fetch_tunnel_ids() -> List[str]:
-            from prime_tunnel.core.client import TunnelClient
-
-            client = TunnelClient()
+            client = _tunnel_client_class()()
             try:
                 scoped_team_id = team_id
                 if scoped_team_id is None:
@@ -400,9 +404,7 @@ async def fetch_tunnel_ids() -> List[str]:
     if labels:
 
         async def validate_label_scope() -> None:
-            from prime_tunnel.core.client import TunnelClient
-
-            client = TunnelClient()
+            client = _tunnel_client_class()()
             try:
                 scoped_user_id = client.config.user_id if only_mine else None
                 scoped_team_id = team_id if team_id is not None else client.config.team_id
@@ -454,9 +456,7 @@ async def validate_label_scope() -> None:
         return
 
     async def delete_tunnels() -> tuple[List[str], List[dict], List[dict]]:
-        from prime_tunnel.core.client import TunnelClient
-
-        client = TunnelClient()
+        client = _tunnel_client_class()()
         succeeded: List[str] = []
         not_found: List[dict] = []
         failed: List[dict] = []