nixoTestLib: split setting up function for port-forwarding and setuping up flake

pkgs/testing/nixos_test_lib/nix_setup.py

@@ -71,3 +71,26 @@ def setup_nix_in_nix(closure_info: str | None) -> None:
                         text=True,
                         check=True,
                     )
+
+
+def prepare_test_flake(
+    temp_dir: str, clan_core_for_checks: str, closure_info: str
+) -> str:
+    """Set up Nix store and copy test flake to temporary directory
+
+    Args:
+        temp_dir: Temporary directory
+        clan_core_for_checks: Path to clan-core-for-checks
+        closure_info: Path to closure info for Nix store setup
+
+    Returns:
+        Path to the test flake directory
+    """
+    # Set up Nix store
+    setup_nix_in_nix(closure_info)
+
+    # Copy test flake
+    flake_dir = Path(temp_dir) / "test-flake"
+    subprocess.run(["cp", "-r", clan_core_for_checks, flake_dir], check=True)  # noqa: S603, S607
+    subprocess.run(["chmod", "-R", "+w", flake_dir], check=True)  # noqa: S603, S607
+    return str(flake_dir)

pkgs/testing/nixos_test_lib/ssh.py

@@ -1,34 +1,31 @@
 """SSH and test setup utilities"""
 
-import subprocess
 from pathlib import Path
 from typing import NamedTuple
 
-from .nix_setup import setup_nix_in_nix
 from .port import find_free_port, setup_port_forwarding
 
 
-class TestEnvironment(NamedTuple):
+class SSHConnection(NamedTuple):
     host_port: int
     ssh_key: str
-    flake_dir: str
 
 
-def setup_test_environment(
+def setup_ssh_connection(
     target,
     temp_dir: str,
-    closure_info: str,
     assets_ssh_privkey: str,
-    clan_core_for_checks: str,
-) -> TestEnvironment:
-    """Set up common test environment including SSH, port forwarding, and flake setup
+) -> SSHConnection:
+    """Set up SSH connection with port forwarding to test VM
+
+    Args:
+        target: Test VM target
+        temp_dir: Temporary directory for SSH key
+        assets_ssh_privkey: Path to SSH private key asset
 
     Returns:
-        TestEnvironment with host_port, ssh_key, and flake_dir
+        SSHConnection with host_port and ssh_key path
     """
-    # Run setup function
-    setup_nix_in_nix(closure_info)
-
     host_port = find_free_port()
     target.wait_for_unit("sshd.service")
     target.wait_for_open_port(22)
@@ -40,9 +37,4 @@ def setup_test_environment(
         f.write(src.read())
     ssh_key.chmod(0o600)
 
-    # Copy test flake to temp directory
-    flake_dir = Path(temp_dir) / "test-flake"
-    subprocess.run(["cp", "-r", clan_core_for_checks, flake_dir], check=True)  # noqa: S603, S607
-    subprocess.run(["chmod", "-R", "+w", flake_dir], check=True)  # noqa: S603, S607
-
-    return TestEnvironment(host_port, str(ssh_key), str(flake_dir))
+    return SSHConnection(host_port, str(ssh_key))