dev/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

nixos-rebuild-ng: run -> run_wrapper, handle encode errors and add extra_env

Browse Source
This commit is contained in:
Thiago Kenji Okada 2024-11-17 19:10:03 +00:00
parent e47b17e239
commit 31e9e8c0aa
5 changed files with 159 additions and 80 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py
View File

@ -13,7 +13,7 @@ from .models import (
NRError,
Profile,
)
from .process import run
from .process import run_wrapper
from .utils import Args, dict_to_flags, info
FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"]
@ -22,7 +22,7 @@ FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"]
def edit(flake: Flake | None, **flake_flags: Args) -> None:
"Try to find and open NixOS configuration file in editor."
if flake:
run(
run_wrapper(
[
"nix",
*FLAKE_FLAGS,
@ -38,9 +38,8 @@ def edit(flake: Flake | None, **flake_flags: Args) -> None:
raise NRError("'edit' does not support extra Nix flags")
nixos_config = Path(
os.getenv("NIXOS_CONFIG")
or run(
or run_wrapper(
["nix-instantiate", "--find-file", "nixos-config"],
text=True,
stdout=PIPE,
check=False,
).stdout.strip()
@ -50,18 +49,17 @@ def edit(flake: Flake | None, **flake_flags: Args) -> None:
nixos_config /= "default.nix"
if nixos_config.exists():
run([os.getenv("EDITOR", "nano"), nixos_config], check=False)
run_wrapper([os.getenv("EDITOR", "nano"), nixos_config], check=False)
else:
raise NRError("cannot find NixOS config file")
def find_file(file: str, **nix_flags: Args) -> Path | None:
"Find classic Nixpkgs location."
r = run(
r = run_wrapper(
["nix-instantiate", "--find-file", file, *dict_to_flags(nix_flags)],
stdout=PIPE,
check=False,
text=True,
)
if r.returncode:
return None
@ -81,17 +79,19 @@ def get_nixpkgs_rev(nixpkgs_path: Path | None) -> str | None:
return None
# Get current revision
r = run(
r = run_wrapper(
["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
check=False,
stdout=PIPE,
text=True,
)
rev = r.stdout.strip()
if rev:
# Check if repo is dirty
if run(["git", "-C", nixpkgs_path, "diff", "--quiet"], check=False).returncode:
if run_wrapper(
["git", "-C", nixpkgs_path, "diff", "--quiet"],
check=False,
).returncode:
rev += "M"
return f".git.{rev}"
else:
@ -142,10 +142,9 @@ def get_generations(profile: Profile, lock_profile: bool = False) -> list[Genera
# 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(
r = run_wrapper(
["nix-env", "-p", profile.path, "--list-generations"],
text=True,
stdout=True,
stdout=PIPE,
check=True,
)
for line in r.stdout.splitlines():
@ -185,11 +184,10 @@ def list_generations(profile: Profile) -> list[GenerationJson]:
s.name for s in (generation_path / "specialisation").glob("*") if s.is_dir()
]
try:
configuration_revision = run(
configuration_revision = run_wrapper(
[generation_path / "sw/bin/nixos-version", "--configuration-revision"],
capture_output=True,
check=True,
text=True,
).stdout.strip()
except (CalledProcessError, IOError):
configuration_revision = "Unknown"
@ -233,7 +231,7 @@ def nixos_build(
else:
run_args = ["nix-build", "<nixpkgs/nixos>", "--attr", attr]
run_args += dict_to_flags(nix_flags)
r = run(run_args, check=True, text=True, stdout=PIPE)
r = run_wrapper(run_args, check=True, stdout=PIPE)
return Path(r.stdout.strip())
@ -254,13 +252,13 @@ def nixos_build_flake(
f"{flake}.config.system.build.{attr}",
]
run_args += dict_to_flags(flake_flags)
r = run(run_args, check=True, text=True, stdout=PIPE)
r = run_wrapper(run_args, check=True, stdout=PIPE)
return Path(r.stdout.strip())
def rollback(profile: Profile) -> Path:
"Rollback Nix profile, like one created by `nixos-rebuild switch`."
run(["nix-env", "--rollback", "-p", profile.path], check=True)
run_wrapper(["nix-env", "--rollback", "-p", profile.path], check=True)
# Rollback config PATH is the own profile
return profile.path
@ -281,7 +279,9 @@ def rollback_temporary_profile(profile: Profile) -> Path | None:
def set_profile(profile: Profile, path_to_config: Path, sudo: bool) -> None:
"Set a path as the current active Nix profile."
run(["nix-env", "-p", profile.path, "--set", path_to_config], check=True, sudo=sudo)
run_wrapper(
["nix-env", "-p", profile.path, "--set", path_to_config], check=True, sudo=sudo
)
def switch_to_configuration(
@ -306,12 +306,9 @@ def switch_to_configuration(
if not path_to_config.exists():
raise NRError(f"specialisation not found: {specialisation}")
run(
run_wrapper(
[path_to_config / "bin/switch-to-configuration", str(action)],
env={
"NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0",
"LOCALE_ARCHIVE": os.getenv("LOCALE_ARCHIVE", ""),
},
extra_env={"NIXOS_INSTALL_BOOTLOADER": "1" if install_bootloader else "0"},
check=True,
sudo=sudo,
)
@ -329,4 +326,4 @@ def upgrade_channels(all: bool = False) -> None:
or channel_path.name == "nixos"
or (channel_path / ".update-on-nixos-rebuild").exists()
):
run(["nix-channel", "--update", channel_path.name], check=False)
run_wrapper(["nix-channel", "--update", channel_path.name], check=False)

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

@ -1,20 +1,52 @@
from __future__ import annotations
import os
import subprocess
from typing import Any, Sequence
from typing import Sequence, TypedDict, Unpack
from .models import SSH
def run(
args: Sequence[str | bytes | os.PathLike[Any]],
# make `check` explicit so we always know if the code is aborting on errors
check: bool,
# Not exhaustive, but we can always extend it later.
class RunKwargs(TypedDict, total=False):
capture_output: bool
stderr: int | None
stdin: int | None
stdout: int | None
def run_wrapper(
args: Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]],
*,
check: bool, # make it explicit so we always know if the code is handling errors
env: dict[str, str] | None = None, # replaces the current environment
extra_env: dict[str, str] | None = None, # appends to the current environment
remote: SSH | None = None,
sudo: bool = False,
**kwargs: Any,
) -> subprocess.CompletedProcess[Any]:
if sudo:
args = ["sudo"] + list(args)
**kwargs: Unpack[RunKwargs],
) -> subprocess.CompletedProcess[str]:
"Wrapper around `subprocess.run` that supports extra functionality."
if remote:
args = ["ssh", *remote.opts, remote.host, "--"] + list(args)
return subprocess.run(args, check=check, **kwargs)
assert env is None, "'env' can't be used with 'remote'"
if extra_env:
extra_env_args = [f"{env}={value}" for env, value in extra_env.items()]
args = ["env", *extra_env_args, *args]
if sudo:
args = ["sudo", *args]
args = ["ssh", *remote.opts, remote.host, "--", *args]
else:
if extra_env:
env = (env or os.environ) | extra_env
if sudo:
args = ["sudo", *args]
return subprocess.run(
args,
check=check,
env=env,
# Hope nobody is using NixOS with non-UTF8 encodings, but "surrogateescape"
# should still work in those systems.
text=True,
errors="surrogateescape",
**kwargs,
)

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

@ -2,7 +2,7 @@ import textwrap
from pathlib import Path
from subprocess import PIPE, CompletedProcess
from typing import Any
from unittest.mock import call, patch
from unittest.mock import ANY, call, patch
import pytest
@ -10,10 +10,7 @@ import nixos_rebuild as nr
from .helpers import get_qualified_name
@pytest.fixture(autouse=True)
def setup(monkeypatch: Any) -> None:
monkeypatch.setenv("LOCALE_ARCHIVE", "/locale")
DEFAULT_RUN_KWARGS = {"text": True, "errors": "surrogateescape", "env": ANY}
def test_parse_args() -> None:
@ -70,6 +67,7 @@ def test_parse_args() -> None:
assert r2.attr == "bar"
@patch.dict(nr.process.os.environ, {}, clear=True)
@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
@patch(get_qualified_name(nr.nix.shutil.which), autospec=True, return_value="/bin/git")
def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> None:
@ -99,17 +97,18 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
["nix-instantiate", "--find-file", "nixpkgs", "-vvv"],
stdout=PIPE,
check=False,
text=True,
**DEFAULT_RUN_KWARGS,
),
call(
["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
check=False,
stdout=PIPE,
text=True,
**DEFAULT_RUN_KWARGS,
),
call(
["git", "-C", nixpkgs_path, "diff", "--quiet"],
check=False,
**DEFAULT_RUN_KWARGS,
),
call(
[
@ -121,8 +120,8 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
"-vvv",
],
check=True,
text=True,
stdout=PIPE,
**DEFAULT_RUN_KWARGS,
),
call(
[
@ -133,16 +132,18 @@ def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> Non
config_path,
],
check=True,
**DEFAULT_RUN_KWARGS,
),
call(
[config_path / "bin/switch-to-configuration", "boot"],
env={"NIXOS_INSTALL_BOOTLOADER": "0", "LOCALE_ARCHIVE": "/locale"},
check=True,
**(DEFAULT_RUN_KWARGS | {"env": {"NIXOS_INSTALL_BOOTLOADER": "0"}}),
),
]
)
@patch.dict(nr.process.os.environ, {}, clear=True)
@patch(get_qualified_name(nr.process.subprocess.run), autospec=True)
def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
config_path = tmp_path / "test"
@ -183,8 +184,8 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
"-v",
],
check=True,
text=True,
stdout=PIPE,
**DEFAULT_RUN_KWARGS,
),
call(
[
@ -196,11 +197,12 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None:
config_path,
],
check=True,
**DEFAULT_RUN_KWARGS,
),
call(
["sudo", config_path / "bin/switch-to-configuration", "switch"],
env={"NIXOS_INSTALL_BOOTLOADER": "1", "LOCALE_ARCHIVE": "/locale"},
check=True,
**(DEFAULT_RUN_KWARGS | {"env": {"NIXOS_INSTALL_BOOTLOADER": "1"}}),
),
]
)
@ -225,14 +227,15 @@ def test_execute_switch_rollback(mock_run: Any, tmp_path: Path) -> None:
Path("/nix/var/nix/profiles/system"),
],
check=True,
**DEFAULT_RUN_KWARGS,
),
call(
[
Path("/nix/var/nix/profiles/system/bin/switch-to-configuration"),
"switch",
],
env={"NIXOS_INSTALL_BOOTLOADER": "1", "LOCALE_ARCHIVE": "/locale"},
check=True,
**DEFAULT_RUN_KWARGS,
),
]
)
@ -281,9 +284,9 @@ def test_execute_test_rollback(
Path("/nix/var/nix/profiles/system-profiles/foo"),
"--list-generations",
],
text=True,
stdout=True,
check=True,
stdout=PIPE,
**DEFAULT_RUN_KWARGS,
),
call(
[
@ -292,8 +295,8 @@ def test_execute_test_rollback(
),
"test",
],
env={"NIXOS_INSTALL_BOOTLOADER": "0", "LOCALE_ARCHIVE": "/locale"},
check=True,
**DEFAULT_RUN_KWARGS,
),
]
)

40
pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py
View File

@ -12,7 +12,7 @@ import nixos_rebuild.nix as n
from .helpers import get_qualified_name
@patch(get_qualified_name(n.run, n), autospec=True)
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_edit(mock_run: Any, monkeypatch: Any, tmpdir: Any) -> None:
# Flake
flake = m.Flake.parse(".#attr")
@ -49,7 +49,7 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
path = Path("/path/to/nix")
with patch(
get_qualified_name(n.run, n),
get_qualified_name(n.run_wrapper, n),
autospec=True,
side_effect=[CompletedProcess([], 0, "")],
) as mock_run:
@ -58,7 +58,6 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
["git", "-C", path, "rev-parse", "--short", "HEAD"],
check=False,
stdout=PIPE,
text=True,
)
expected_calls = [
@ -66,7 +65,6 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
["git", "-C", path, "rev-parse", "--short", "HEAD"],
check=False,
stdout=PIPE,
text=True,
),
call(
["git", "-C", path, "diff", "--quiet"],
@ -75,7 +73,7 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
]
with patch(
get_qualified_name(n.run, n),
get_qualified_name(n.run_wrapper, n),
autospec=True,
side_effect=[
CompletedProcess([], 0, "0f7c82403fd6"),
@ -86,7 +84,7 @@ def test_get_nixpkgs_rev(mock_which: Any) -> None:
mock_run.assert_has_calls(expected_calls)
with patch(
get_qualified_name(n.run, n),
get_qualified_name(n.run_wrapper, n),
autospec=True,
side_effect=[
CompletedProcess([], 0, "0f7c82403fd6"),
@ -118,7 +116,7 @@ def test_get_generations_from_nix_store(tmp_path: Path) -> None:
@patch(
get_qualified_name(n.run, n),
get_qualified_name(n.run_wrapper, n),
autospec=True,
return_value=CompletedProcess(
[],
@ -183,7 +181,7 @@ def test_list_generations(mock_get_generations: Any, tmp_path: Path) -> None:
@patch(
get_qualified_name(n.run, n),
get_qualified_name(n.run_wrapper, n),
autospec=True,
return_value=CompletedProcess([], 0, stdout=" \n/path/to/file\n "),
)
@ -209,13 +207,12 @@ def test_nixos_build_flake(mock_run: Any) -> None:
"foo",
],
check=True,
text=True,
stdout=PIPE,
)
@patch(
get_qualified_name(n.run, n),
get_qualified_name(n.run_wrapper, n),
autospec=True,
return_value=CompletedProcess([], 0, stdout=" \n/path/to/file\n "),
)
@ -224,7 +221,6 @@ def test_nixos_build(mock_run: Any, monkeypatch: Any) -> None:
mock_run.assert_called_with(
["nix-build", "<nixpkgs/nixos>", "--attr", "attr", "--nix-flag", "foo"],
check=True,
text=True,
stdout=PIPE,
)
@ -232,7 +228,6 @@ def test_nixos_build(mock_run: Any, monkeypatch: Any) -> None:
mock_run.assert_called_with(
["nix-build", "file", "--attr", "preAttr.attr"],
check=True,
text=True,
stdout=PIPE,
)
@ -240,7 +235,6 @@ def test_nixos_build(mock_run: Any, monkeypatch: Any) -> None:
mock_run.assert_called_with(
["nix-build", "file", "--attr", "attr", "--no-out-link"],
check=True,
text=True,
stdout=PIPE,
)
@ -248,12 +242,11 @@ def test_nixos_build(mock_run: Any, monkeypatch: Any) -> None:
mock_run.assert_called_with(
["nix-build", "default.nix", "--attr", "preAttr.attr", "--keep-going"],
check=True,
text=True,
stdout=PIPE,
)
@patch(get_qualified_name(n.run, n), autospec=True)
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_rollback(mock_run: Any, tmp_path: Path) -> None:
path = tmp_path / "test"
path.touch()
@ -269,7 +262,7 @@ def test_rollback_temporary_profile(tmp_path: Path) -> None:
path.touch()
profile = m.Profile("system", path)
with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
mock_run.return_value = CompletedProcess(
[],
0,
@ -288,12 +281,12 @@ def test_rollback_temporary_profile(tmp_path: Path) -> None:
== path.parent / "foo-2083-link"
)
with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
mock_run.return_value = CompletedProcess([], 0, stdout="")
assert n.rollback_temporary_profile(profile) is None
@patch(get_qualified_name(n.run, n), autospec=True)
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_set_profile(mock_run: Any) -> None:
profile_path = Path("/path/to/profile")
config_path = Path("/path/to/config")
@ -306,7 +299,7 @@ def test_set_profile(mock_run: Any) -> None:
)
@patch(get_qualified_name(n.run, n), autospec=True)
@patch(get_qualified_name(n.run_wrapper, n), autospec=True)
def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
profile_path = Path("/path/to/profile")
config_path = Path("/path/to/config")
@ -323,7 +316,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
)
mock_run.assert_called_with(
[profile_path / "bin/switch-to-configuration", "switch"],
env={"NIXOS_INSTALL_BOOTLOADER": "0", "LOCALE_ARCHIVE": ""},
extra_env={"NIXOS_INSTALL_BOOTLOADER": "0"},
check=True,
sudo=False,
)
@ -342,6 +335,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
with monkeypatch.context() as mp:
mp.setenv("LOCALE_ARCHIVE", "/path/to/locale")
mp.setenv("PATH", "/path/to/bin")
mp.setattr(Path, Path.exists.__name__, lambda self: True)
n.switch_to_configuration(
@ -356,7 +350,7 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
config_path / "specialisation/special/bin/switch-to-configuration",
"test",
],
env={"NIXOS_INSTALL_BOOTLOADER": "1", "LOCALE_ARCHIVE": "/path/to/locale"},
extra_env={"NIXOS_INSTALL_BOOTLOADER": "1"},
check=True,
sudo=True,
)
@ -372,11 +366,11 @@ def test_switch_to_configuration(mock_run: Any, monkeypatch: Any) -> None:
],
)
def test_upgrade_channels(mock_glob: Any) -> None:
with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
n.upgrade_channels(False)
mock_run.assert_called_with(["nix-channel", "--update", "nixos"], check=False)
with patch(get_qualified_name(n.run, n), autospec=True) as mock_run:
with patch(get_qualified_name(n.run_wrapper, n), autospec=True) as mock_run:
n.upgrade_channels(True)
mock_run.assert_has_calls(
[

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

@ -1,6 +1,8 @@
from typing import Any
from unittest.mock import patch
import pytest
import nixos_rebuild.models as m
import nixos_rebuild.process as p
@ -9,13 +11,34 @@ from .helpers import get_qualified_name
@patch(get_qualified_name(p.subprocess.run))
def test_run(mock_run: Any) -> None:
p.run(["test", "--with", "flags"], check=True)
mock_run.assert_called_with(["test", "--with", "flags"], check=True)
p.run_wrapper(["test", "--with", "flags"], check=True)
mock_run.assert_called_with(
["test", "--with", "flags"],
check=True,
text=True,
errors="surrogateescape",
env=None,
)
p.run(["test", "--with", "flags"], check=False, sudo=True)
mock_run.assert_called_with(["sudo", "test", "--with", "flags"], check=False)
with patch.dict(p.os.environ, {"PATH": "/path/to/bin"}, clear=True):
p.run_wrapper(
["test", "--with", "flags"],
check=False,
sudo=True,
extra_env={"FOO": "bar"},
)
mock_run.assert_called_with(
["sudo", "test", "--with", "flags"],
check=False,
text=True,
errors="surrogateescape",
env={
"PATH": "/path/to/bin",
"FOO": "bar",
},
)
p.run(
p.run_wrapper(
["test", "--with", "flags"],
check=True,
remote=m.SSH("user@localhost", ["--ssh", "opt"]),
@ -23,12 +46,29 @@ def test_run(mock_run: Any) -> None:
mock_run.assert_called_with(
["ssh", "--ssh", "opt", "user@localhost", "--", "test", "--with", "flags"],
check=True,
text=True,
errors="surrogateescape",
env=None,
)
p.run(
p.run_wrapper(
["test", "--with", "flags"],
check=False,
env={"FOO": "bar"},
)
mock_run.assert_called_with(
["test", "--with", "flags"],
check=False,
env={"FOO": "bar"},
text=True,
errors="surrogateescape",
)
p.run_wrapper(
["test", "--with", "flags"],
check=True,
sudo=True,
extra_env={"FOO": "bar"},
remote=m.SSH("user@localhost", ["--ssh", "opt"]),
)
mock_run.assert_called_with(
@ -39,9 +79,22 @@ def test_run(mock_run: Any) -> None:
"user@localhost",
"--",
"sudo",
"env",
"FOO=bar",
"test",
"--with",
"flags",
],
check=True,
env=None,
text=True,
errors="surrogateescape",
)
with pytest.raises(AssertionError):
p.run_wrapper(
["test", "--with", "flags"],
check=False,
env={"foo": "bar"},
remote=m.SSH("user@localhost", []),
)