diff options
author | Patrick Hilhorst <git@hilhorst.be> | 2022-01-01 22:35:20 +0100 |
---|---|---|
committer | Patrick Hilhorst <git@hilhorst.be> | 2022-01-01 23:17:32 +0100 |
commit | 4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2 (patch) | |
tree | 30e14d078a28988a7234c66b41fb93e3bd0f1b6b /nixos/lib/test-driver/test_driver | |
parent | 69856d9ba78905337407136f48012c23962871e7 (diff) | |
download | nixpkgs-4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2.tar nixpkgs-4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2.tar.gz nixpkgs-4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2.tar.bz2 nixpkgs-4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2.tar.lz nixpkgs-4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2.tar.xz nixpkgs-4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2.tar.zst nixpkgs-4e1556ed4d43da1f930b3fcf0fc20d827a34f3d2.zip |
nixos/test-driver: add polling_condition
Diffstat (limited to 'nixos/lib/test-driver/test_driver')
-rw-r--r-- | nixos/lib/test-driver/test_driver/driver.py | 40 | ||||
-rw-r--r-- | nixos/lib/test-driver/test_driver/machine.py | 6 | ||||
-rw-r--r-- | nixos/lib/test-driver/test_driver/polling_condition.py | 90 |
3 files changed, 135 insertions, 1 deletions
diff --git a/nixos/lib/test-driver/test_driver/driver.py b/nixos/lib/test-driver/test_driver/driver.py index f3af98537ad..e22f9ee7a75 100644 --- a/nixos/lib/test-driver/test_driver/driver.py +++ b/nixos/lib/test-driver/test_driver/driver.py @@ -1,12 +1,13 @@ from contextlib import contextmanager from pathlib import Path -from typing import Any, Dict, Iterator, List +from typing import Any, Dict, Iterator, List, Union, Optional, Callable, ContextManager import os import tempfile from test_driver.logger import rootlog from test_driver.machine import Machine, NixStartScript, retry from test_driver.vlan import VLan +from test_driver.polling_condition import PollingCondition class Driver: @@ -16,6 +17,7 @@ class Driver: tests: str vlans: List[VLan] machines: List[Machine] + polling_conditions: List[Callable] def __init__( self, @@ -36,12 +38,15 @@ class Driver: for s in scripts: yield NixStartScript(s) + self.polling_conditions = [] + self.machines = [ Machine( start_command=cmd, keep_vm_state=keep_vm_state, name=cmd.machine_name, tmp_dir=tmp_dir, + fail_early=self.fail_early, ) for cmd in cmd(start_scripts) ] @@ -84,6 +89,7 @@ class Driver: retry=retry, serial_stdout_off=self.serial_stdout_off, serial_stdout_on=self.serial_stdout_on, + polling_condition=self.polling_condition, Machine=Machine, # for typing ) machine_symbols = {m.name: m for m in self.machines} @@ -159,3 +165,35 @@ class Driver: def serial_stdout_off(self) -> None: rootlog._print_serial_logs = False + + def fail_early(self) -> bool: + return any(not f() for f in self.polling_conditions) + + def polling_condition( + self, + fun_: Optional[Callable] = None, + *, + seconds_interval: float = 2.0, + description: Optional[str] = None, + ) -> Union[Callable[[Callable], ContextManager], ContextManager]: + driver = self + + class Poll: + def __init__(self, fun: Callable): + self.condition = PollingCondition( + fun, + seconds_interval, + description, + ).check + + def __enter__(self) -> None: + driver.polling_conditions.append(self.condition) + + def __exit__(self, a, b, c) -> None: # type: ignore + res = driver.polling_conditions.pop() + assert res is self.condition + + if fun_ is None: + return Poll + else: + return Poll(fun_) diff --git a/nixos/lib/test-driver/test_driver/machine.py b/nixos/lib/test-driver/test_driver/machine.py index b3dbe5126fc..dbf9fd24486 100644 --- a/nixos/lib/test-driver/test_driver/machine.py +++ b/nixos/lib/test-driver/test_driver/machine.py @@ -17,6 +17,7 @@ import threading import time from test_driver.logger import rootlog +from test_driver.polling_condition import PollingCondition, coopmulti CHAR_TO_KEY = { "A": "shift-a", @@ -318,6 +319,7 @@ class Machine: # Store last serial console lines for use # of wait_for_console_text last_lines: Queue = Queue() + fail_early: Callable def __repr__(self) -> str: return f"<Machine '{self.name}'>" @@ -329,12 +331,14 @@ class Machine: name: str = "machine", keep_vm_state: bool = False, allow_reboot: bool = False, + fail_early: Callable = lambda: False, ) -> None: self.tmp_dir = tmp_dir self.keep_vm_state = keep_vm_state self.allow_reboot = allow_reboot self.name = name self.start_command = start_command + self.fail_early = fail_early # set up directories self.shared_dir = self.tmp_dir / "shared-xchg" @@ -405,6 +409,7 @@ class Machine: break return answer + @coopmulti def send_monitor_command(self, command: str) -> str: with self.nested("sending monitor command: {}".format(command)): message = ("{}\n".format(command)).encode() @@ -506,6 +511,7 @@ class Machine: break return "".join(output_buffer) + @coopmulti def execute( self, command: str, check_return: bool = True, timeout: Optional[int] = 900 ) -> Tuple[int, str]: diff --git a/nixos/lib/test-driver/test_driver/polling_condition.py b/nixos/lib/test-driver/test_driver/polling_condition.py new file mode 100644 index 00000000000..f38dea71376 --- /dev/null +++ b/nixos/lib/test-driver/test_driver/polling_condition.py @@ -0,0 +1,90 @@ +from typing import Callable, Optional, Any, List, Dict +from functools import wraps + +import time + +from .logger import rootlog + + +class PollingConditionFailed(Exception): + pass + + +def coopmulti(fun: Callable, *, machine: Any = None) -> Callable: + assert not (fun is None and machine is None) + + def inner(fun_: Callable) -> Any: + @wraps(fun_) + def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any: + this_machine = args[0] if machine is None else machine + + if this_machine.fail_early(): # type: ignore + raise PollingConditionFailed("Action interrupted early...") + + return fun_(*args, **kwargs) + + return wrapper + + if fun is None: + return inner + else: + return inner(fun) + + +class PollingCondition: + condition: Callable[[], bool] + seconds_interval: float + description: Optional[str] + + last_called: float + entered: bool + + def __init__( + self, + condition: Callable[[], Optional[bool]], + seconds_interval: float = 2.0, + description: Optional[str] = None, + ): + self.condition = condition # type: ignore + self.seconds_interval = seconds_interval + + if description is None: + self.description = condition.__doc__ + else: + self.description = str(description) + + self.last_called = float("-inf") + self.entered = False + + def check(self) -> bool: + if self.entered or not self.overdue: + return True + + with self, rootlog.nested(self.nested_message): + rootlog.info(f"Time since last: {time.monotonic() - self.last_called:.2f}s") + try: + res = self.condition() # type: ignore + except Exception: + res = False + res = res is None or res + rootlog.info(f"Polling condition {'succeeded' if res else 'failed'}") + return res + + @property + def nested_message(self) -> str: + nested_message = ["Checking polling condition"] + if self.description is not None: + nested_message.append(repr(self.description)) + + return " ".join(nested_message) + + @property + def overdue(self) -> bool: + return self.last_called + self.seconds_interval < time.monotonic() + + def __enter__(self) -> None: + self.entered = True + + def __exit__(self, exc_type, exc_value, traceback) -> None: # type: ignore + self.entered = False + self.last_called = time.monotonic() |