diff --git a/goth/runner/step.py b/goth/runner/step.py index 513adcf6..5c33e994 100644 --- a/goth/runner/step.py +++ b/goth/runner/step.py @@ -5,6 +5,8 @@ import time from typing import Optional, TYPE_CHECKING +import colors + from goth.runner.exceptions import StepTimeoutError if TYPE_CHECKING: @@ -19,14 +21,18 @@ def step(default_timeout: float = 10.0): def decorator(func): @functools.wraps(func) - async def wrapper(self: "Probe", *args, timeout: Optional[float] = None): + async def wrapper( + self: "Probe", *args, timeout: Optional[float] = None, **kwargs + ): timeout = timeout if timeout is not None else default_timeout step_name = f"{self.name}.{func.__name__}(timeout={timeout})" start_time = time.time() logger.info("Running step '%s'", step_name) try: - result = await asyncio.wait_for(func(self, *args), timeout=timeout) + result = await asyncio.wait_for( + func(self, *args, **kwargs), timeout=timeout + ) self.runner.check_assertion_errors() step_time = time.time() - start_time logger.debug( @@ -37,12 +43,16 @@ async def wrapper(self: "Probe", *args, timeout: Optional[float] = None): ) except asyncio.TimeoutError: step_time = time.time() - start_time - logger.error("Step '%s' timed out after %.1f s", step_name, step_time) + logger.error( + colors.red("Step '%s' timed out after %.1f s", style="bold"), + step_name, + step_time, + ) raise StepTimeoutError(step_name, step_time) except Exception as exc: step_time = time.time() - start_time logger.error( - "Step '%s' raised %s in %.1f", + colors.red("Step '%s' raised %s in %.1f", style="bold"), step_name, exc.__class__.__name__, step_time,