From 0212452e27161b583dfee83b5bf2f52141a379bb Mon Sep 17 00:00:00 2001 From: Thiago Kenji Okada Date: Sun, 15 Dec 2024 12:38:55 +0000 Subject: [PATCH] nixos-rebuild-ng: split get_generations in 2 functions --- .../nixos-rebuild-ng/src/nixos_rebuild/nix.py | 117 +++++++++--------- .../ni/nixos-rebuild-ng/src/tests/test_nix.py | 58 ++++++--- 2 files changed, 98 insertions(+), 77 deletions(-) diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py index c7033808394b..c1f2e5d5c711 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py @@ -266,68 +266,72 @@ def get_nixpkgs_rev(nixpkgs_path: Path | None) -> str | None: return None -def _parse_generation_from_nix_store(path: Path, profile: Profile) -> Generation: - entry_id = path.name.split("-")[1] - current = path.name == profile.path.readlink().name - timestamp = datetime.fromtimestamp(path.stat().st_ctime).strftime( - "%Y-%m-%d %H:%M:%S" - ) - - return Generation( - id=int(entry_id), - timestamp=timestamp, - current=current, - ) - - -def _parse_generation_from_nix_env(line: str) -> Generation: - parts = line.split() - - entry_id = parts[0] - timestamp = f"{parts[1]} {parts[2]}" - current = "(current)" in parts - - return Generation( - id=int(entry_id), - timestamp=timestamp, - current=current, - ) - - -def get_generations( - profile: Profile, - target_host: Remote | None = None, - using_nix_env: bool = False, - sudo: bool = False, -) -> list[Generation]: +def get_generations(profile: Profile) -> list[Generation]: """Get all NixOS generations from profile. Includes generation ID (e.g.: 1, 2), timestamp (e.g.: when it was created) and if this is the current active profile or not. - - If `lock_profile = True` this command will need root to run successfully. """ if not profile.path.exists(): raise NRError(f"no profile '{profile.name}' found") - result = [] - if using_nix_env: - # Using `nix-env --list-generations` needs root to lock the profile - # TODO: do we actually need to lock profile for e.g.: rollback? - # https://github.com/NixOS/nix/issues/5144 - r = run_wrapper( - ["nix-env", "-p", profile.path, "--list-generations"], - stdout=PIPE, - remote=target_host, - sudo=sudo, + def parse_path(path: Path, profile: Profile) -> Generation: + entry_id = path.name.split("-")[1] + current = path.name == profile.path.readlink().name + timestamp = datetime.fromtimestamp(path.stat().st_ctime).strftime( + "%Y-%m-%d %H:%M:%S" ) - for line in r.stdout.splitlines(): - result.append(_parse_generation_from_nix_env(line)) - else: - assert not target_host, "target_host is not supported when using_nix_env=False" - for p in profile.path.parent.glob("system-*-link"): - result.append(_parse_generation_from_nix_store(p, profile)) - return sorted(result, key=lambda d: d.id) + + return Generation( + id=int(entry_id), + timestamp=timestamp, + current=current, + ) + + return sorted( + [parse_path(p, profile) for p in profile.path.parent.glob("system-*-link")], + key=lambda d: d.id, + ) + + +def get_generations_from_nix_env( + profile: Profile, + target_host: Remote | None = None, + sudo: bool = False, +) -> list[Generation]: + """Get all NixOS generations from profile with nix-env. Needs root. + + Includes generation ID (e.g.: 1, 2), timestamp (e.g.: when it was created) + and if this is the current active profile or not. + """ + if not profile.path.exists(): + raise NRError(f"no profile '{profile.name}' found") + + # Using `nix-env --list-generations` needs root to lock the profile + r = run_wrapper( + ["nix-env", "-p", profile.path, "--list-generations"], + stdout=PIPE, + remote=target_host, + sudo=sudo, + ) + + def parse_line(line: str) -> Generation: + parts = line.split() + + entry_id = parts[0] + timestamp = f"{parts[1]} {parts[2]}" + current = "(current)" in parts + + return Generation( + id=int(entry_id), + timestamp=timestamp, + current=current, + ) + + return sorted( + [parse_line(line) for line in r.stdout.splitlines()], + key=lambda d: d.id, + ) def list_generations(profile: Profile) -> list[GenerationJson]: @@ -437,11 +441,8 @@ def rollback_temporary_profile( sudo: bool, ) -> Path | None: "Rollback a temporary Nix profile, like one created by `nixos-rebuild test`." - generations = get_generations( - profile, - target_host=target_host, - using_nix_env=True, - sudo=sudo, + generations = get_generations_from_nix_env( + profile, target_host=target_host, sudo=sudo ) previous_gen_id = None for generation in generations: diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py index 3cee46045863..34f47029c103 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py @@ -327,7 +327,7 @@ def test_get_nixpkgs_rev() -> None: mock_run.assert_has_calls(expected_calls) -def test_get_generations_from_nix_store(tmp_path: Path) -> None: +def test_get_generations(tmp_path: Path) -> None: nixos_path = tmp_path / "nixos-system" nixos_path.mkdir() @@ -337,20 +337,17 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None: (tmp_path / "system-3-link").symlink_to(nixos_path) (tmp_path / "system-2-link").symlink_to(nixos_path) - assert n.get_generations( - m.Profile("system", tmp_path / "system"), - using_nix_env=False, - ) == [ + assert n.get_generations(m.Profile("system", tmp_path / "system")) == [ m.Generation(id=1, current=False, timestamp=ANY), m.Generation(id=2, current=True, timestamp=ANY), m.Generation(id=3, current=False, timestamp=ANY), ] -@patch( - get_qualified_name(n.run_wrapper, n), - autospec=True, - return_value=CompletedProcess( +def test_get_generations_from_nix_env(tmp_path: Path) -> None: + path = tmp_path / "test" + path.touch() + return_value = CompletedProcess( [], 0, stdout=textwrap.dedent("""\ @@ -358,17 +355,40 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None: 2083 2024-11-07 22:59:41 2084 2024-11-07 23:54:17 (current) """), - ), -) -def test_get_generations_from_nix_env(mock_run: Any, tmp_path: Path) -> None: - path = tmp_path / "test" - path.touch() + ) - assert n.get_generations(m.Profile("system", path), using_nix_env=True) == [ - m.Generation(id=2082, current=False, timestamp="2024-11-07 22:58:56"), - m.Generation(id=2083, current=False, timestamp="2024-11-07 22:59:41"), - m.Generation(id=2084, current=True, timestamp="2024-11-07 23:54:17"), - ] + with patch( + get_qualified_name(n.run_wrapper, n), autospec=True, return_value=return_value + ) as mock_run: + assert n.get_generations_from_nix_env(m.Profile("system", path)) == [ + m.Generation(id=2082, current=False, timestamp="2024-11-07 22:58:56"), + m.Generation(id=2083, current=False, timestamp="2024-11-07 22:59:41"), + m.Generation(id=2084, current=True, timestamp="2024-11-07 23:54:17"), + ] + mock_run.assert_called_with( + ["nix-env", "-p", path, "--list-generations"], + stdout=PIPE, + remote=None, + sudo=False, + ) + + remote = m.Remote("user@host", [], "password") + with patch( + get_qualified_name(n.run_wrapper, n), autospec=True, return_value=return_value + ) as mock_run: + assert n.get_generations_from_nix_env( + m.Profile("system", path), remote, True + ) == [ + m.Generation(id=2082, current=False, timestamp="2024-11-07 22:58:56"), + m.Generation(id=2083, current=False, timestamp="2024-11-07 22:59:41"), + m.Generation(id=2084, current=True, timestamp="2024-11-07 23:54:17"), + ] + mock_run.assert_called_with( + ["nix-env", "-p", path, "--list-generations"], + stdout=PIPE, + remote=remote, + sudo=True, + ) @patch(