nixos-rebuild-ng: accept protocol in Flake.parse()

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

@@ -181,7 +181,7 @@ def parse_args(
     }
 
     if args.help or args.action is None:
-        if "@withShellFiles@" == "true":
+        if "@withShellFiles@" == "true":  # type: ignore
             r = run(["man", "8", "@executable@"], check=False)
             parser.exit(r.returncode)
         else:
@@ -298,7 +298,7 @@ def execute(argv: list[str]) -> None:
     # Re-exec to a newer version of the script before building to ensure we get
     # the latest fixes
     if (
-        "@withReexec@" == "true"
+        "@withReexec@" == "true"  # type: ignore
         and can_run
         and not args.fast
         and not os.environ.get("_NIXOS_REBUILD_REEXEC")

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

@@ -60,7 +60,7 @@ class BuildAttr:
 
 @dataclass(frozen=True)
 class Flake:
-    path: Path
+    path: Path | str
     attr: str
     _re: ClassVar = re.compile(r"^(?P<path>[^\#]*)\#?(?P<attr>[^\#\"]*)$")
 
@@ -81,7 +81,11 @@ class Flake:
         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"}"
-        return cls(Path(m.group("path")), nixos_attr)
+        path = m.group("path")
+        if ":" in path:
+            return cls(path, nixos_attr)
+        else:
+            return cls(Path(path), nixos_attr)
 
     @classmethod
     def from_arg(cls, flake_arg: Any, target_host: Remote | None) -> Self | None:

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

@@ -474,6 +474,46 @@ def test_execute_build(mock_run: Any, tmp_path: Path) -> None:
     )
 
 
+@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
+def test_execute_test_flake(mock_run: Any, tmp_path: Path) -> None:
+    config_path = tmp_path / "test"
+    config_path.touch()
+    mock_run.side_effect = [
+        # nixos_build_flake
+        CompletedProcess([], 0, str(config_path)),
+        # switch_to_configuration
+        CompletedProcess([], 0),
+    ]
+
+    nr.execute(
+        ["nixos-rebuild", "test", "--flake", "github:user/repo#hostname", "--fast"]
+    )
+
+    assert mock_run.call_count == 2
+    mock_run.assert_has_calls(
+        [
+            call(
+                [
+                    "nix",
+                    "--extra-experimental-features",
+                    "nix-command flakes",
+                    "build",
+                    "--print-out-paths",
+                    "github:user/repo#nixosConfigurations.hostname.config.system.build.toplevel",
+                ],
+                check=True,
+                stdout=PIPE,
+                **DEFAULT_RUN_KWARGS,
+            ),
+            call(
+                [config_path / "bin/switch-to-configuration", "test"],
+                check=True,
+                **DEFAULT_RUN_KWARGS,
+            ),
+        ]
+    )
+
+
 @patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
 @patch(get_qualified_name(nr.nix.Path.exists, nr.nix), autospec=True, return_value=True)
 @patch(get_qualified_name(nr.nix.Path.mkdir, nr.nix), autospec=True)

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

@@ -43,6 +43,12 @@ def test_flake_parse() -> None:
     assert m.Flake.parse(".#attr") == m.Flake(Path("."), "nixosConfigurations.attr")
     assert m.Flake.parse("#attr") == m.Flake(Path("."), "nixosConfigurations.attr")
     assert m.Flake.parse(".") == m.Flake(Path("."), "nixosConfigurations.default")
+    assert m.Flake.parse("path:/to/flake#attr") == m.Flake(
+        "path:/to/flake", "nixosConfigurations.attr"
+    )
+    assert m.Flake.parse("github:user/repo/branch") == m.Flake(
+        "github:user/repo/branch", "nixosConfigurations.default"
+    )
 
 
 def test_flake_to_attr() -> None: