nixos-rebuild-ng: move temporary directory to process

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py

@@ -6,7 +6,6 @@ import os
 import sys
 from pathlib import Path
 from subprocess import run
-from tempfile import TemporaryDirectory
 from typing import assert_never
 
 from . import nix
@@ -164,16 +163,11 @@ def execute(argv: list[str]) -> None:
     flake_build_flags = common_build_flags | vars(args_groups["flake_build_flags"])
     copy_flags = common_flags | vars(args_groups["copy_flags"])
 
-    # Will be cleaned up on exit automatically.
-    tmpdir = TemporaryDirectory(prefix="nixos-rebuild.")
-    tmpdir_path = Path(tmpdir.name)
-    atexit.register(cleanup_ssh, tmpdir_path)
+    atexit.register(cleanup_ssh)
 
     profile = Profile.from_arg(args.profile_name)
-    build_host = Remote.from_arg(
-        args.build_host, False, tmpdir_path, validate_opts=False
-    )
-    target_host = Remote.from_arg(args.target_host, args.ask_sudo_password, tmpdir_path)
+    build_host = Remote.from_arg(args.build_host, False, validate_opts=False)
+    target_host = Remote.from_arg(args.target_host, args.ask_sudo_password)
     build_attr = BuildAttr.from_arg(args.attr, args.file)
     flake = Flake.from_arg(args.flake, target_host)
     action = Action(args.action)
@@ -200,8 +194,7 @@ def execute(argv: list[str]) -> None:
                 argv[0],
                 new,
             )
-            cleanup_ssh(tmpdir_path)
-            tmpdir.cleanup()
+            cleanup_ssh()
             os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
 
     if args.upgrade or args.upgrade_all:

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py

@@ -5,10 +5,22 @@ import subprocess
 from dataclasses import dataclass
 from getpass import getpass
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from typing import Self, Sequence, TypedDict, Unpack
 
 logger = logging.getLogger(__name__)
 
+TMPDIR = TemporaryDirectory(prefix="nixos-rebuild.")
+TMPDIR_PATH = Path(TMPDIR.name)
+SSH_DEFAULT_OPTS = [
+    "-o",
+    "ControlMaster=auto",
+    "-o",
+    f"ControlPath={TMPDIR_PATH / "ssh-%n"}",
+    "-o",
+    "ControlPersist=60",
+]
+
 
 @dataclass(frozen=True)
 class Remote:
@@ -21,7 +33,6 @@ class Remote:
         cls,
         host: str | None,
         ask_sudo_password: bool | None,
-        tmp_dir: Path,
         validate_opts: bool = True,
     ) -> Self | None:
         if not host:
@@ -30,15 +41,6 @@ class Remote:
         opts = os.getenv("NIX_SSHOPTS", "").split()
         if validate_opts:
             cls._validate_opts(opts, ask_sudo_password)
-        opts += [
-            # SSH ControlMaster flags, allow for faster re-connection
-            "-o",
-            "ControlMaster=auto",
-            "-o",
-            f"ControlPath={tmp_dir / "ssh-%n"}",
-            "-o",
-            "ControlPersist=60",
-        ]
         sudo_password = None
         if ask_sudo_password:
             sudo_password = getpass(f"[sudo] password for {host}: ")
@@ -66,12 +68,13 @@ class RunKwargs(TypedDict, total=False):
     stdout: int | None
 
 
-def cleanup_ssh(tmp_dir: Path) -> None:
+def cleanup_ssh() -> None:
     "Close SSH ControlMaster connection."
-    for ctrl in tmp_dir.glob("ssh-*"):
+    for ctrl in TMPDIR_PATH.glob("ssh-*"):
         subprocess.run(
             ["ssh", "-o", f"ControlPath={ctrl}", "-O", "exit", "dummyhost"], check=False
         )
+    TMPDIR.cleanup()
 
 
 def run_wrapper(
@@ -99,6 +102,7 @@ def run_wrapper(
         args = [
             "ssh",
             *remote.opts,
+            *SSH_DEFAULT_OPTS,
             remote.host,
             "--",
             # sadly SSH just join all remaining parameters, expanding glob and

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py

@@ -218,9 +218,9 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
 
 @patch.dict(nr.process.os.environ, {}, clear=True)
 @patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
-@patch(get_qualified_name(nr.TemporaryDirectory, nr))  # can't autospec
+@patch(get_qualified_name(nr.cleanup_ssh, nr), autospec=True)
 def test_execute_nix_switch_flake_target_host(
-    mock_tmpdir: Any,
+    mock_cleanup_ssh: Any,
     mock_run: Any,
     tmp_path: Path,
 ) -> None:
@@ -236,7 +236,6 @@ def test_execute_nix_switch_flake_target_host(
         # switch_to_configuration
         CompletedProcess([], 0),
     ]
-    mock_tmpdir.return_value.name = "/tmp/test"
 
     nr.execute(
         [
@@ -276,12 +275,7 @@ def test_execute_nix_switch_flake_target_host(
             call(
                 [
                     "ssh",
-                    "-o",
-                    "ControlMaster=auto",
-                    "-o",
-                    "ControlPath=/tmp/test/ssh-%n",
-                    "-o",
-                    "ControlPersist=60",
+                    *nr.process.SSH_DEFAULT_OPTS,
                     "user@localhost",
                     "--",
                     f"sudo nix-env -p /nix/var/nix/profiles/system --set {config_path}",
@@ -292,12 +286,7 @@ def test_execute_nix_switch_flake_target_host(
             call(
                 [
                     "ssh",
-                    "-o",
-                    "ControlMaster=auto",
-                    "-o",
-                    "ControlPath=/tmp/test/ssh-%n",
-                    "-o",
-                    "ControlPersist=60",
+                    *nr.process.SSH_DEFAULT_OPTS,
                     "user@localhost",
                     "--",
                     f"sudo env NIXOS_INSTALL_BOOTLOADER=0 {config_path / 'bin/switch-to-configuration'} switch",
@@ -311,9 +300,9 @@ def test_execute_nix_switch_flake_target_host(
 
 @patch.dict(nr.process.os.environ, {}, clear=True)
 @patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
-@patch(get_qualified_name(nr.TemporaryDirectory, nr))  # can't autospec
+@patch(get_qualified_name(nr.cleanup_ssh, nr), autospec=True)
 def test_execute_nix_switch_flake_build_host(
-    mock_tmpdir: Any,
+    mock_cleanup_ssh: Any,
     mock_run: Any,
     tmp_path: Path,
 ) -> None:
@@ -331,7 +320,6 @@ def test_execute_nix_switch_flake_build_host(
         # switch_to_configuration
         CompletedProcess([], 0),
     ]
-    mock_tmpdir.return_value.name = "/tmp/test"
 
     nr.execute(
         [
@@ -369,12 +357,7 @@ def test_execute_nix_switch_flake_build_host(
             call(
                 [
                     "ssh",
-                    "-o",
-                    "ControlMaster=auto",
-                    "-o",
-                    "ControlPath=/tmp/test/ssh-%n",
-                    "-o",
-                    "ControlPersist=60",
+                    *nr.process.SSH_DEFAULT_OPTS,
                     "user@localhost",
                     "--",
                     f"nix --extra-experimental-features 'nix-command flakes' build '{config_path}^*' --print-out-paths",

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py

@@ -1,6 +1,5 @@
-from pathlib import Path
 from typing import Any
-from unittest.mock import call, patch
+from unittest.mock import patch
 
 import nixos_rebuild.models as m
 import nixos_rebuild.process as p
@@ -8,29 +7,7 @@ import nixos_rebuild.process as p
 from .helpers import get_qualified_name
 
 
-@patch(get_qualified_name(p.subprocess.run))
-def test_cleanup_ssh(mock_run: Any, tmp_path: Path) -> None:
-    (tmp_path / "ssh-conn").touch()
-
-    p.cleanup_ssh(tmp_path)
-    mock_run.assert_has_calls(
-        [
-            call(
-                [
-                    "ssh",
-                    "-o",
-                    f"ControlPath={tmp_path}/ssh-conn",
-                    "-O",
-                    "exit",
-                    "dummyhost",
-                ],
-                check=False,
-            )
-        ]
-    )
-
-
-@patch(get_qualified_name(p.subprocess.run))
+@patch(get_qualified_name(p.subprocess.run), autospec=True)
 def test_run(mock_run: Any) -> None:
     p.run_wrapper(["test", "--with", "flags"], check=True)
     mock_run.assert_called_with(
@@ -67,7 +44,15 @@ def test_run(mock_run: Any) -> None:
         remote=m.Remote("user@localhost", ["--ssh", "opt"], "password"),
     )
     mock_run.assert_called_with(
-        ["ssh", "--ssh", "opt", "user@localhost", "--", "test --with 'some flags'"],
+        [
+            "ssh",
+            "--ssh",
+            "opt",
+            *p.SSH_DEFAULT_OPTS,
+            "user@localhost",
+            "--",
+            "test --with 'some flags'",
+        ],
         check=True,
         text=True,
         errors="surrogateescape",
@@ -87,6 +72,7 @@ def test_run(mock_run: Any) -> None:
             "ssh",
             "--ssh",
             "opt",
+            *p.SSH_DEFAULT_OPTS,
             "user@localhost",
             "--",
             "sudo --prompt= --stdin env FOO=bar test --with flags",
@@ -99,38 +85,20 @@ def test_run(mock_run: Any) -> None:
     )
 
 
-def test_remote_from_name(monkeypatch: Any, tmpdir: Path) -> None:
+def test_remote_from_name(monkeypatch: Any) -> None:
     monkeypatch.setenv("NIX_SSHOPTS", "")
-    assert m.Remote.from_arg("user@localhost", None, tmpdir) == m.Remote(
+    assert m.Remote.from_arg("user@localhost", None, False) == m.Remote(
         "user@localhost",
-        opts=[
-            "-o",
-            "ControlMaster=auto",
-            "-o",
-            f"ControlPath={tmpdir / "ssh-%n"}",
-            "-o",
-            "ControlPersist=60",
-        ],
+        opts=[],
         sudo_password=None,
     )
 
     # get_qualified_name doesn't work because getpass is aliased to another
     # function
     with patch(f"{p.__name__}.getpass", return_value="password"):
-        monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar")
-        assert m.Remote.from_arg("user@localhost", True, tmpdir) == m.Remote(
+        monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar -t")
+        assert m.Remote.from_arg("user@localhost", True, True) == m.Remote(
             "user@localhost",
-            opts=[
-                "-f",
-                "foo",
-                "-b",
-                "bar",
-                "-o",
-                "ControlMaster=auto",
-                "-o",
-                f"ControlPath={tmpdir / "ssh-%n"}",
-                "-o",
-                "ControlPersist=60",
-            ],
+            opts=["-f", "foo", "-b", "bar", "-t"],
             sudo_password="password",
         )