Files
lix/tests/functional2/testlib/fixtures/nix.py
T
Commentator2.0 01985e5add functional2: use loggers
Use logger in favor over print statment.
This is explicitly supported and encuraged by pytest, which also allows
for capturing logs separate from stdout calls, which is handy for when
e.g. lix code calls out to stdout to keep those differentiated from test
output

Change-Id: Ib88565a1663da3b77ca6b95f8edf644eafb4a99d
2025-05-10 21:13:45 +02:00

244 lines
8.0 KiB
Python

import dataclasses
import json
import logging
import os
import subprocess
from functools import partialmethod
from pathlib import Path
from typing import Any, AnyStr
from collections.abc import Callable
import pytest
from functional2.testlib.terminal_code_eater import eat_terminal_codes
logger = logging.getLogger(__name__)
@dataclasses.dataclass
class CommandResult:
cmd: list[str]
rc: int
"""Return code"""
stderr: bytes
"""Outputted stderr"""
stdout: bytes
"""Outputted stdout"""
def ok(self) -> "CommandResult":
if self.rc != 0:
logger.debug("stdout: %s", self.stderr_s)
logger.debug("stderr: %s", self.stderr_s)
raise subprocess.CalledProcessError(
returncode=self.rc, cmd=self.cmd, stderr=self.stderr, output=self.stdout
)
return self
def expect(self, rc: int) -> "CommandResult":
if self.rc != rc:
logger.debug("stdout: %s", self.stderr_s)
logger.debug("stderr: %s", self.stderr_s)
raise subprocess.CalledProcessError(
returncode=self.rc, cmd=self.cmd, stderr=self.stderr, output=self.stdout
)
return self
@property
def stdout_s(self) -> str:
"""Command stdout as str"""
return self.stdout.decode("utf-8", errors="replace")
@property
def stderr_s(self) -> str:
"""Command stderr as str"""
return self.stderr.decode("utf-8", errors="replace")
@property
def stdout_plain(self) -> str:
"""Command stderr as str with terminal escape sequences eaten and whitespace stripped"""
return eat_terminal_codes(self.stdout).decode("utf-8", errors="replace").strip()
@property
def stderr_plain(self) -> str:
"""Command stderr as str with terminal escape sequences eaten and whitespace stripped"""
return eat_terminal_codes(self.stderr).decode("utf-8", errors="replace").strip()
def json(self) -> Any:
self.ok()
return json.loads(self.stdout)
@dataclasses.dataclass
class NixSettings:
"""Settings for invoking Nix"""
experimental_features: set[str] | None = None
store: str | None = None
"""
The store to operate on (may be a path or other thing, see nix help-stores).
Note that this can be set to the test's store directory if you want to use
/nix/store paths inside that test rather than NIX_STORE_DIR renaming
/nix/store to some unstable name (assuming that no builds are invoked).
"""
nix_store_dir: Path | None = None
"""
Alternative name to use for /nix/store: breaks all references and NAR imports if
set, but does allow builds in tests (since builds do not require chroots if
the store is relocated).
"""
def feature(self, *names: str) -> "NixSettings":
self.experimental_features = (self.experimental_features or set()) | set(names)
return self
def to_config(self) -> str:
config = ""
def serialise(value: Any) -> str:
if type(value) in {str, int}:
return str(value)
if type(value) in {list, set}:
return " ".join(str(e) for e in value)
msg = f"Value is unsupported in nix config: {value!r}"
raise ValueError(msg)
def field_may(name: str, value: Any, serializer: Callable[[Any], str] = serialise):
nonlocal config
if value is not None:
config += f"{name} = {serializer(value)}\n"
field_may("experimental-features", self.experimental_features)
field_may("store", self.store)
assert (
self.store or self.nix_store_dir
), "Failing to set either nix_store_dir or store will cause accidental use of the system store."
return config
def to_env_overlay(self) -> dict[str, str]:
ret = {"NIX_CONFIG": self.to_config()}
if self.nix_store_dir:
ret["NIX_STORE_DIR"] = str(self.nix_store_dir)
return ret
@dataclasses.dataclass
class Command:
argv: list[str]
env: dict[str, str] = dataclasses.field(default_factory=dict)
stdin: bytes | None = None
cwd: Path | None = None
def with_env(self, **kwargs) -> "Command":
self.env.update(kwargs)
return self
def with_stdin(self, stdin: bytes) -> "Command":
self.stdin = stdin
return self
def run(self) -> CommandResult:
proc = subprocess.Popen(
self.argv,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
stdin=subprocess.PIPE if self.stdin else subprocess.DEVNULL,
cwd=self.cwd,
env=self.env,
)
(stdout, stderr) = proc.communicate(input=self.stdin)
rc = proc.returncode
return CommandResult(cmd=self.argv, rc=rc, stdout=stdout, stderr=stderr)
@dataclasses.dataclass
class NixCommand(Command):
settings: NixSettings = dataclasses.field(default_factory=NixSettings)
def apply_nix_config(self):
self.env.update(self.settings.to_env_overlay())
def run(self) -> CommandResult:
self.apply_nix_config()
return super().run()
@dataclasses.dataclass
class Nix:
test_root: Path
def hermetic_env(self) -> dict[str, Path]:
# mirroring vars-and-functions.sh
home = self.test_root / "test-home"
home.mkdir(parents=True, exist_ok=True)
return {
"NIX_LOCALSTATE_DIR": self.test_root / "var",
"NIX_LOG_DIR": self.test_root / "var/log/nix",
"NIX_STATE_DIR": self.test_root / "var/nix",
"NIX_CONF_DIR": self.test_root / "etc",
"NIX_DAEMON_SOCKET_PATH": self.test_root / "daemon-socket",
"NIX_USER_CONF_FILES": "",
"HOME": home,
}
def make_env(self) -> dict[AnyStr, AnyStr]:
# We conservatively assume that people might want to successfully get
# some env through to the subprocess, so we override whatever is in the
# global env.
d = os.environ.copy()
d.update(self.hermetic_env())
return d
def cmd(self, argv: list[str]) -> Command:
return Command(argv=argv, cwd=self.test_root, env=self.make_env())
def settings(self, allow_builds: bool = False) -> NixSettings:
"""
Parameters:
- allow_builds: relocate the Nix store so that builds work (however, makes store paths non-reproducible across test runs!)
"""
settings = NixSettings()
store_path = self.test_root / "store"
if allow_builds:
settings.nix_store_dir = store_path
else:
settings.store = str(store_path)
return settings
def nix_cmd(self, argv: list[str], flake: bool = False) -> NixCommand:
"""
Constructs a NixCommand with the appropriate settings.
"""
settings = self.settings()
if flake:
settings.feature("nix-command", "flakes")
return NixCommand(argv=argv, cwd=self.test_root, env=self.make_env(), settings=settings)
def nix(self, cmd: list[str], nix_exe: str = "nix", flake: bool = False) -> NixCommand:
return self.nix_cmd([nix_exe, *cmd], flake=flake)
nix_build = partialmethod(nix, nix_exe="nix-build")
nix_shell = partialmethod(nix, nix_exe="nix-shell")
nix_store = partialmethod(nix, nix_exe="nix-store")
nix_env = partialmethod(nix, nix_exe="nix-env")
nix_instantiate = partialmethod(nix, nix_exe="nix-instantiate")
nix_channel = partialmethod(nix, nix_exe="nix-channel")
nix_prefetch_url = partialmethod(nix, nix_exe="nix-prefetch-url")
def eval(self, expr: str, settings: NixSettings | None = None) -> CommandResult:
if settings is None:
settings = self.settings()
# clone due to reference-shenanigans
settings = dataclasses.replace(settings).feature("nix-command")
cmd = self.nix(["eval", "--json", "--expr", expr])
cmd.settings = settings
return cmd.run()
@pytest.fixture
def nix(tmp_path: Path) -> Nix:
return Nix(tmp_path)