From 7830f000c57bb616b178a6a8eaef9659938ca7ea Mon Sep 17 00:00:00 2001 From: Patrick Hilhorst Date: Sun, 2 Jan 2022 22:20:04 +0100 Subject: nixos/test-driver: simplify coopmulti --- .../test-driver/test_driver/polling_condition.py | 24 +++++++--------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/nixos/lib/test-driver/test_driver/polling_condition.py b/nixos/lib/test-driver/test_driver/polling_condition.py index f38dea71376..fe064b1f830 100644 --- a/nixos/lib/test-driver/test_driver/polling_condition.py +++ b/nixos/lib/test-driver/test_driver/polling_condition.py @@ -10,25 +10,15 @@ class PollingConditionFailed(Exception): pass -def coopmulti(fun: Callable, *, machine: Any = None) -> Callable: - assert not (fun is None and machine is None) +def coopmulti(fun: Callable) -> Callable: + @wraps(fun) + def wrapper(machine: Any, *args: List[Any], **kwargs: Dict[str, Any]) -> Any: + if machine.fail_early(): # type: ignore + raise PollingConditionFailed("Test interrupted early...") - 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 + return fun(machine, *args, **kwargs) - 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) + return wrapper class PollingCondition: -- cgit 1.4.1