dynix/tests/harmonia/test-script.py
2026-02-18 16:17:06 +01:00

115 lines
4.2 KiB
Python

import functools
from pathlib import Path
import shlex
import textwrap
import tomllib
from typing import Any, cast, TYPE_CHECKING
from beartype import beartype
from test_driver.machine import Machine
from test_driver.errors import RequestedAssertionFailed
if TYPE_CHECKING:
global machine
machine = cast(Machine, ...)
assert machine.shell is not None
ls = "eza -lah --color=always --group-directories-first"
indent = functools.partial(textwrap.indent, prefix=' ')
@beartype
def run_log(machine: Machine, *commands: str, timeout: int | None = 60) -> str:
output = ""
for command in commands:
with machine.nested(f"must succeed: {command}"):
(status, out) = machine.execute(f"{command} | tee /dev/stderr", timeout=timeout)
if status != 0:
machine.log(f"output: {out}")
raise RequestedAssertionFailed(
f"command `{command}` failed (exit code {status})",
)
output += out
return output
@beartype
def get_config_file() -> dict[str, Any]:
machine.wait_for_unit("harmonia.service")
pid = int(machine.get_unit_property("harmonia.service", "MainPID"))
env_lines: list[str] = machine.succeed(f"cat /proc/{pid}/environ").replace("\0", "\n").splitlines()
pairs: list[list[str]] = [line.split("=", maxsplit=1) for line in env_lines]
env = dict(pairs)
config_file = Path(env["CONFIG_FILE"])
machine.log(f"copying from VM: {config_file=}")
machine.copy_from_vm(config_file.as_posix())
config_file_path = machine.out_dir / config_file.name
with open(config_file_path, "rb") as f:
config_data = tomllib.load(f)
config_file_path.unlink()
return config_data
@beartype
def dynix_append(option: str, value: Any):
machine.succeed(f'''
dynix append {shlex.quote(option)} {shlex.quote(str(value))}
'''.strip())
@beartype
def do_apply():
expr = textwrap.dedent("""
(import <nixpkgs/nixos> { }).config.dynamicism.applyDynamicConfiguration { }
""").strip()
machine.succeed(rf"""
nix run --show-trace --log-format raw-with-logs --impure -E {shlex.quote(expr)}
""".strip())
machine.wait_for_unit("default.target")
machine.wait_for_unit("install-dynix.service")
dynix_out = machine.succeed("dynix --version")
assert "dynix" in dynix_out, f"dynix not in {dynix_out=}"
# Config should have our initial values.
config_toml = get_config_file()
assert int(config_toml['workers']) == 4, f"{config_toml['workers']=} != 4"
assert int(config_toml['max_connection_rate']) == 256, f"{config_toml['max_connection_rate']=} != 256"
with machine.nested("must succeed: initial nixos-rebuild switch"):
machine.succeed("env PAGER= nixos-rebuild switch --log-format raw-with-logs --fallback")
# Config should not have changed.
config_toml = get_config_file()
assert int(config_toml['workers']) == 4, f"{config_toml['workers']=} != 4"
assert int(config_toml['max_connection_rate']) == 256, f"{config_toml['max_connection_rate']=} != 256"
new_workers = 20
dynix_append("services.harmonia.settings.workers", new_workers)
do_apply()
# Workers, but not max connection rate, should have changed.
config_toml = get_config_file()
assert int(config_toml['workers']) == new_workers, f"{config_toml['workers']=} != {new_workers}"
assert int(config_toml['max_connection_rate']) == 256, f"{config_toml['max_connection_rate']=} != 256"
new_max_connection_rate = 100
dynix_append("services.harmonia.settings.max_connection_rate", new_max_connection_rate)
do_apply()
# Max connection rate should have changed, and workers should be the same as before.
config_toml = get_config_file()
assert int(config_toml['max_connection_rate']) == new_max_connection_rate, f"{config_toml['max_connection_rate']=} != {new_max_connection_rate}"
assert int(config_toml['workers']) == new_workers, f"{config_toml['workers']=} != {new_workers}"
# And this should set everything back.
machine.succeed("env PAGER= nixos-rebuild switch --log-format raw-with-logs --fallback")
machine.wait_for_unit("harmonia.service")
config_toml = get_config_file()
assert int(config_toml['max_connection_rate']) == 256, f'{config_toml["max_connection_rate"]=} != 256'
assert int(config_toml['workers']) == 4, f'{config_toml["workers"]=} != 4'