From 8481fa781e03ca06d084f353a9ce2ff355a9e95f Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Thu, 19 Jun 2025 20:48:48 +0100
Subject: [PATCH] nixos-rebuild-ng: simplify `Flake.parse()` code

---
 .../src/nixos_rebuild/models.py               | 45 +++++++++----------
 .../nixos-rebuild-ng/src/tests/test_models.py | 25 ++++++++---
 2 files changed, 41 insertions(+), 29 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py
index 0d67bc81aed9..f8439461b6bd 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py
@@ -4,7 +4,7 @@ import subprocess
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Any, Callable, ClassVar, Self, TypedDict, override
+from typing import Any, ClassVar, Self, TypedDict, override
 
 from .process import Remote, run_wrapper
 
@@ -100,6 +100,20 @@ def discover_closest_flake(location: Path) -> Path | None:
     return None
 
 
+def get_hostname(target_host: Remote | None) -> str | None:
+    if target_host:
+        try:
+            return run_wrapper(
+                ["uname", "-n"],
+                capture_output=True,
+                remote=target_host,
+            ).stdout.strip()
+        except (AttributeError, subprocess.CalledProcessError):
+            return None
+    else:
+        return platform.node()
+
+
 @dataclass(frozen=True)
 class Flake:
     path: Path | str
@@ -114,15 +128,13 @@ class Flake:
         return f"{self.path}#{self.attr}"
 
     @classmethod
-    def parse(
-        cls,
-        flake_str: str,
-        hostname_fn: Callable[[], str | None] = lambda: None,
-    ) -> Self:
+    def parse(cls, flake_str: str, target_host: Remote | None = None) -> Self:
         m = cls._re.match(flake_str)
         assert m is not None, f"got no matches for {flake_str}"
         attr = m.group("attr")
-        nixos_attr = f'nixosConfigurations."{attr or hostname_fn() or "default"}"'
+        nixos_attr = (
+            f'nixosConfigurations."{attr or get_hostname(target_host) or "default"}"'
+        )
         path_str = m.group("path")
         if ":" in path_str:
             return cls(path_str, nixos_attr)
@@ -143,24 +155,11 @@ class Flake:
 
     @classmethod
     def from_arg(cls, flake_arg: Any, target_host: Remote | None) -> Self | None:
-        def get_hostname() -> str | None:
-            if target_host:
-                try:
-                    return run_wrapper(
-                        ["uname", "-n"],
-                        stdout=subprocess.PIPE,
-                        remote=target_host,
-                    ).stdout.strip()
-                except (AttributeError, subprocess.CalledProcessError):
-                    return None
-            else:
-                return platform.node()
-
         match flake_arg:
             case str(s):
-                return cls.parse(s, get_hostname)
+                return cls.parse(s, target_host)
             case True:
-                return cls.parse(".", get_hostname)
+                return cls.parse(".", target_host)
             case False:
                 return None
             case _:
@@ -169,7 +168,7 @@ class Flake:
                 if default_path.exists():
                     # It can be a symlink to the actual flake.
                     default_path = default_path.resolve()
-                    return cls.parse(str(default_path.parent), get_hostname)
+                    return cls.parse(str(default_path.parent), target_host)
                 else:
                     return None
 
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_models.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_models.py
index 2cc450e00a36..1922b4682cbf 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_models.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_models.py
@@ -30,16 +30,23 @@ def test_build_attr_to_attr() -> None:
     )
 
 
-def test_flake_parse(tmpdir: Path, monkeypatch: MonkeyPatch) -> None:
+@patch("platform.node", autospec=True, return_value="hostname")
+def test_flake_parse(mock_node: Mock, tmpdir: Path, monkeypatch: MonkeyPatch) -> None:
     assert m.Flake.parse("/path/to/flake#attr") == m.Flake(
         Path("/path/to/flake"), 'nixosConfigurations."attr"'
     )
-    assert m.Flake.parse("/path/ to /flake", lambda: "hostname") == m.Flake(
+    assert m.Flake.parse("/path/ to /flake") == m.Flake(
         Path("/path/ to /flake"), 'nixosConfigurations."hostname"'
     )
-    assert m.Flake.parse("/path/to/flake", lambda: "hostname") == m.Flake(
-        Path("/path/to/flake"), 'nixosConfigurations."hostname"'
-    )
+    with patch(
+        get_qualified_name(m.run_wrapper, m),
+        autospec=True,
+        return_value=subprocess.CompletedProcess([], 0, stdout="remote\n"),
+    ):
+        target_host = m.Remote("target@remote", [], None)
+        assert m.Flake.parse("/path/to/flake", target_host) == m.Flake(
+            Path("/path/to/flake"), 'nixosConfigurations."remote"'
+        )
     # change directory to tmpdir
     with monkeypatch.context() as patch_context:
         patch_context.chdir(tmpdir)
@@ -49,10 +56,16 @@ def test_flake_parse(tmpdir: Path, monkeypatch: MonkeyPatch) -> None:
         assert m.Flake.parse("#attr") == m.Flake(
             Path("."), 'nixosConfigurations."attr"'
         )
-        assert m.Flake.parse(".") == m.Flake(Path("."), 'nixosConfigurations."default"')
+        assert m.Flake.parse(".") == m.Flake(
+            Path("."), 'nixosConfigurations."hostname"'
+        )
     assert m.Flake.parse("path:/to/flake#attr") == m.Flake(
         "path:/to/flake", 'nixosConfigurations."attr"'
     )
+
+    # from here on  we should return "default"
+    mock_node.return_value = None
+
     assert m.Flake.parse("github:user/repo/branch") == m.Flake(
         "github:user/repo/branch", 'nixosConfigurations."default"'
     )