move setup_nix_in_nix into nixos_test_lib

checks/installation/flake-module.nix

@@ -168,7 +168,6 @@
             nodes.target = (import ./test-helpers.nix { inherit lib pkgs self; }).target;
             extraPythonPackages = _p: [
               (import ./test-helpers.nix { inherit lib pkgs self; }).nixosTestLib
-              self.legacyPackages.${pkgs.system}.setupNixInNixPythonPackage
             ];
 
             testScript = ''
@@ -227,7 +226,6 @@
             nodes.target = (import ./test-helpers.nix { inherit lib pkgs self; }).target;
             extraPythonPackages = _p: [
               (import ./test-helpers.nix { inherit lib pkgs self; }).nixosTestLib
-              self.legacyPackages.${pkgs.system}.setupNixInNixPythonPackage
             ];
 
             testScript = ''

checks/installation/nixos_test_lib/nix_setup.py

@@ -0,0 +1,70 @@
+"""Nix store setup utilities for VM tests"""
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+
+def setup_nix_in_nix(closure_info: str) -> None:
+    """Set up Nix store inside test environment
+
+    Args:
+        closure_info: Path to closure info directory containing store-paths file
+    """
+    tmpdir = os.environ.get("TMPDIR", "/tmp")
+
+    # Remove NIX_REMOTE if present (we don't have any nix daemon running)
+    if "NIX_REMOTE" in os.environ:
+        del os.environ["NIX_REMOTE"]
+
+    # Set NIX_CONFIG globally to disable substituters for speed
+    os.environ["NIX_CONFIG"] = "substituters = \ntrusted-public-keys = "
+
+    # Set up environment variables for test environment
+    os.environ["HOME"] = tmpdir
+    os.environ["NIX_STATE_DIR"] = f"{tmpdir}/nix"
+    os.environ["NIX_CONF_DIR"] = f"{tmpdir}/etc"
+    os.environ["IN_NIX_SANDBOX"] = "1"
+    os.environ["CLAN_TEST_STORE"] = f"{tmpdir}/store"
+    os.environ["LOCK_NIX"] = f"{tmpdir}/nix_lock"
+
+    # Create necessary directories
+    Path(f"{tmpdir}/nix").mkdir(parents=True, exist_ok=True)
+    Path(f"{tmpdir}/etc").mkdir(parents=True, exist_ok=True)
+    Path(f"{tmpdir}/store").mkdir(parents=True, exist_ok=True)
+
+    # Set up Nix store if closure info is provided
+    if closure_info and os.path.exists(closure_info):
+        store_paths_file = os.path.join(closure_info, "store-paths")
+        if os.path.exists(store_paths_file):
+            with open(store_paths_file) as f:
+                store_paths = f.read().strip().split("\n")
+
+            # Copy store paths to test store
+            for store_path in store_paths:
+                if store_path.strip():
+                    dest_path = f"{tmpdir}/store{store_path}"
+                    if not os.path.exists(dest_path):
+                        # Create parent directories
+                        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+                        # Copy the store path
+                        if os.path.isdir(store_path):
+                            shutil.copytree(store_path, dest_path, dirs_exist_ok=True)
+                        else:
+                            shutil.copy2(store_path, dest_path)
+
+            # Load Nix database
+            registration_file = os.path.join(closure_info, "registration")
+            if os.path.exists(registration_file):
+                env = os.environ.copy()
+                env["NIX_REMOTE"] = f"local?store={tmpdir}/store"
+
+                with open(registration_file) as f:
+                    subprocess.run(
+                        ["nix-store", "--load-db"],
+                        input=f.read(),
+                        text=True,
+                        env=env,
+                        check=True,
+                    )

checks/installation/nixos_test_lib/ssh.py

@@ -4,7 +4,7 @@ import os
 import subprocess
 from typing import NamedTuple
 
-from .environment import setup_nix_environment
+from .nix_setup import setup_nix_in_nix
 from .port import find_free_port, setup_port_forwarding
 
 
@@ -26,12 +26,8 @@ def setup_test_environment(
     Returns:
         TestEnvironment with host_port, ssh_key, and flake_dir
     """
-    from setup_nix_in_nix import setup_nix_in_nix  # type: ignore[import-untyped]
-
-    setup_nix_environment(temp_dir, closure_info)
-
     # Run setup function
-    setup_nix_in_nix()
+    setup_nix_in_nix(closure_info)
 
     host_port = find_free_port()
     target.wait_for_unit("sshd.service")

checks/service-dummy-test-from-flake/default.nix

@@ -29,8 +29,8 @@ nixosLib.runTest (
     testScript =
       { nodes, ... }:
       ''
-        from setup_nix_in_nix import setup_nix_in_nix # type: ignore[import-untyped]
+        from nixos_test_lib.nix_setup import setup_nix_in_nix # type: ignore[import-untyped]
-        setup_nix_in_nix()
+        setup_nix_in_nix(None)  # No closure info for this test
 
         def run_clan(cmd: list[str], **kwargs) -> str:
             import subprocess

pkgs/testing/flake-module.nix

@@ -21,57 +21,6 @@
           fi
         '';
 
-        setupNixInNixPythonPackage = pkgs.python3Packages.buildPythonPackage {
-          pname = "setup-nix-in-nix";
-          version = "1.0.0";
-          format = "other";
-
-          dontUnpack = true;
-
-          installPhase = ''
-            mkdir -p $out/${pkgs.python3.sitePackages}
-            cat > $out/${pkgs.python3.sitePackages}/setup_nix_in_nix.py << 'EOF'
-            from os import environ
-            import subprocess
-            from pathlib import Path
-
-            def setup_nix_in_nix():
-                """Set up a Nix store inside the test environment."""
-                environ['HOME'] = environ['TMPDIR']
-                environ['NIX_STATE_DIR'] = environ['TMPDIR'] + '/nix'
-                environ['NIX_CONF_DIR'] = environ['TMPDIR'] + '/etc'
-                environ['IN_NIX_SANDBOX'] = '1'
-                environ['CLAN_TEST_STORE'] = environ['TMPDIR'] + '/store'
-                environ['LOCK_NIX'] = environ['TMPDIR'] + '/nix_lock'
-
-                Path(environ['CLAN_TEST_STORE'] + '/nix/store').mkdir(parents=True, exist_ok=True)
-                Path(environ['CLAN_TEST_STORE'] + '/nix/var/nix/gcroots').mkdir(parents=True, exist_ok=True)
-
-                if 'closureInfo' in environ:
-                    # Read store paths from the closure info file
-                    with open(environ['closureInfo'] + '/store-paths', 'r') as f:
-                        store_paths = f.read().strip().split('\n')
-
-                    # Copy store paths using absolute path to cp
-                    subprocess.run(
-                        ['${pkgs.coreutils}/bin/cp', '--recursive', '--target', environ['CLAN_TEST_STORE'] + '/nix/store'] + store_paths,
-                        check=True
-                    )
-
-                    # Load the nix database using absolute path to nix-store
-                    with open(environ['closureInfo'] + '/registration', 'r') as f:
-                        subprocess.run(
-                            ['${pkgs.nix}/bin/nix-store', '--load-db', '--store', environ['CLAN_TEST_STORE']],
-                            input=f.read(),
-                            text=True,
-                            check=True
-                        )
-            EOF
-            touch $out/${pkgs.python3.sitePackages}/py.typed
-          '';
-
-          doCheck = false;
-        };
       };
     };
 }