Custom failure handling (#353)

* preparations for adding custom failure handling

* implementation and tests custom failure handling with new `RetryTask`

this failure action will cause `grizzly.scenario.iterator` to retry the failed task up to 3 times, with an exponetional random backoff between retries.

if it still fails after 3 retries, the default failure handling will be used (stop user, restart scenario or just continue).

* fixed conditional logging, that wasn't conditional

* handle ServiceBusError for authentication timed-out

* sort variable names when dumping in `--dry-run`

* review feedback fix

example/features/environment.py

@@ -38,4 +38,4 @@ def before_scenario(
     grizzly_before_scenario(context, scenario, *args, **kwargs)
 
     grizzly = cast(GrizzlyContext, context.grizzly)
-    grizzly.scenario.failure_exception = StopUser
+    grizzly.scenario.failure_handling.update({None: StopUser})

grizzly/context.py

@@ -296,7 +296,7 @@ class GrizzlyContextScenario:
     variables: GrizzlyVariables = field(init=False, repr=False, hash=False, default_factory=GrizzlyVariables)
     _tasks: GrizzlyContextTasks = field(init=False, repr=False, hash=False, compare=False)
     validation: GrizzlyContextScenarioValidation = field(init=False, hash=False, compare=False, default_factory=GrizzlyContextScenarioValidation)
-    failure_exception: Optional[type[Exception]] = field(init=False, default=None)
+    failure_handling: dict[type[Exception] | str | None, type[Exception]] = field(init=False, repr=False, hash=False, default_factory=dict)
     orphan_templates: list[str] = field(init=False, repr=False, hash=False, compare=False, default_factory=list)
     _jinja2: Environment = field(init=False, repr=False, default_factory=jinja2_environment_factory)
 

grizzly/exceptions.py

@@ -8,6 +8,7 @@
 from grizzly_extras.exceptions import StopScenario
 
 if TYPE_CHECKING:  # pragma: no cover
+    from grizzly.context import GrizzlyContextScenario
     from grizzly.types.behave import Scenario, Step
 
 
@@ -27,6 +28,11 @@ def __init__(self, message: Optional[str] = None) -> None:
 class RestartScenario(Exception):  # noqa: N818
     pass
 
+
+class RetryTask(Exception):  # noqa: N818
+    pass
+
+
 class AssertionErrors(Exception):  # noqa: N818
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
@@ -45,6 +51,8 @@ def __next__(self) -> AssertionError:
             self._index += 1
             return error
 
+        self._index = 0
+
         raise StopIteration
 
     def __len__(self) -> int:
@@ -79,3 +87,27 @@ def __init__(self, error: Exception) -> None:
 
     def __str__(self) -> str:
         return f'{self.error!s}'
+
+
+def failure_handler(exception: Exception | None, scenario: GrizzlyContextScenario) -> None:
+    # no failure, just return
+    if exception is None:
+        return
+
+    # always raise StopUser when these unhandled exceptions has occured
+    if isinstance(exception, (NotImplementedError, KeyError, IndexError, AttributeError, TypeError, SyntaxError)):
+        raise StopUser from exception
+
+    # custom actions based on failure
+    for failure_type, failure_action in scenario.failure_handling.items():
+        if failure_type is None:
+            continue
+
+        if (isinstance(failure_type, str) and failure_type in str(exception)) or exception.__class__ is failure_type:
+            raise failure_action from exception
+
+    # no match, raise the default if it has been set
+    default_exception = scenario.failure_handling.get(None, None)
+
+    if default_exception is not None:
+        raise default_exception from exception

grizzly/locust.py

@@ -990,7 +990,7 @@ def run(context: Context) -> int:  # noqa: C901, PLR0915, PLR0912
                 for scenario in grizzly.scenarios:
                     logger.info('# %s:', scenario.name)
 
-                    for variable, value in scenario.variables.items():
+                    for variable, value in dict(sorted(scenario.variables.items())).items():
                         if value is None or (isinstance(value, str) and value.lower() == 'none'):
                             continue
                         logger.info('    %s = %s', variable, value)

grizzly/scenarios/iterator.py

@@ -8,6 +8,7 @@
 
 from contextlib import suppress
 from json import dumps as jsondumps
+from random import uniform
 from time import perf_counter
 from typing import TYPE_CHECKING, Any, ClassVar, Optional
 
@@ -17,7 +18,7 @@
 from locust.exception import InterruptTaskSet, RescheduleTask, RescheduleTaskImmediately
 from locust.user.task import LOCUST_STATE_RUNNING, LOCUST_STATE_STOPPING
 
-from grizzly.exceptions import RestartScenario, StopScenario
+from grizzly.exceptions import RestartScenario, RetryTask, StopScenario
 from grizzly.types import RequestType, ScenarioState
 from grizzly.types.locust import StopUser
 
@@ -30,6 +31,14 @@
     from grizzly.users import GrizzlyUser
 
 
+NUMBER_TO_WORD: dict[int, str] = {
+    1: '1st',
+    2: '2nd',
+    3: '3rd',
+    4: '4th',
+}
+
+
 class IteratorScenario(GrizzlyScenario):
     user: GrizzlyUser
 
@@ -120,12 +129,42 @@ def run(self) -> None:  # type: ignore[misc]  # noqa: C901, PLR0912, PLR0915
                     except Exception:
                         step = 'unknown'
 
-                    try:
-                        self.execute_next_task(self.current_task_index + 1, self.task_count, step)
-                    except Exception as e:
-                        if not isinstance(e, StopScenario):
-                            execute_task_logged = True
-                        raise
+                    retries = 0
+                    while True:
+                        try:
+                            self.execute_next_task(self.current_task_index + 1, self.task_count, step)
+                            break
+                        except RetryTask as e:
+                            retries += 1
+
+                            if retries >= 3:
+                                message = f'task {self.current_task_index + 1} of {self.task_count} failed after {retries} retries: {step}'
+                                self.logger.error(message)  # noqa: TRY400
+
+                                default_exception = self.user._scenario.failure_handling.get(None, None)
+
+                                # default failure handling
+                                if default_exception is not None:
+                                    raise default_exception from e
+
+                                break
+
+                            sleep_time = retries * uniform(1.0, 2.0)  # noqa: S311
+                            message = f'task {self.current_task_index + 1} of {self.task_count} will execute a {NUMBER_TO_WORD[retries+1]} time in {sleep_time:.2f} seconds: {step}'
+                            self.logger.warning(message)
+
+                            gsleep(sleep_time)  # random back-off time
+                            self.wait()
+
+                            # step back counter, and re-schedule the same task again
+                            self._task_index -= 1
+                            self.schedule_task(self.get_next_task(), first=True)
+
+                            continue
+                        except Exception as e:
+                            if not isinstance(e, StopScenario):
+                                execute_task_logged = True
+                            raise
                 except RescheduleTaskImmediately:
                     pass
                 except RescheduleTask:

grizzly/steps/scenario/results.py

@@ -32,7 +32,9 @@ def step_results_fail_ratio(context: Context, fail_ratio: int) -> None:
 
     """
     grizzly = cast(GrizzlyContext, context.grizzly)
-    assert grizzly.scenario.failure_exception is None, f"cannot use step 'fail ratio is greater than \"{fail_ratio}\" fail scenario' togheter with 'on failure' steps"
+    assert grizzly.scenario.failure_handling.get(None, None) is None, (
+        f"cannot use step 'fail ratio is greater than \"{fail_ratio}\" fail scenario' together with 'on failure' steps"
+    )
     grizzly.scenario.validation.fail_ratio = fail_ratio / 100.0
 
 

grizzly/steps/scenario/setup.py

@@ -3,6 +3,7 @@
 """
 from __future__ import annotations
 
+from contextlib import suppress
 from typing import Any, Optional, cast
 
 import parse
@@ -12,9 +13,10 @@
 from grizzly.exceptions import RestartScenario
 from grizzly.tasks import GrizzlyTask, RequestTask
 from grizzly.testdata.utils import resolve_variable
-from grizzly.types.behave import Context, given, register_type, then
+from grizzly.types import FailureAction
+from grizzly.types.behave import Context, given, register_type, then, when
 from grizzly.types.locust import StopUser
-from grizzly.utils import has_template
+from grizzly.utils import ModuleLoader, has_template
 from grizzly_extras.text import permutation
 
 
@@ -24,8 +26,33 @@ def parse_iteration_gramatical_number(text: str) -> str:
     return text.strip()
 
 
+def parse_failure_type(value: str) -> type[Exception] | str:
+    result: type[Exception] | str | None = None
+
+    if '.' in value:
+        module_name, value = value.rsplit('.', 1)
+        module_names = [module_name]
+    else:
+        module_names = ['grizzly.exceptions', 'builtins']
+
+    # check if value corresponds to an exception type
+    for module_name in module_names:
+        with suppress(Exception):
+            result = ModuleLoader[Exception].load(module_name, value)
+            break
+
+    # an exception message which should be checked against the string representation
+    # of the exception that was thrown
+    if result is None:
+        result = value
+
+    return result
+
+
 register_type(
     IterationGramaticalNumber=parse_iteration_gramatical_number,
+    FailureType=parse_failure_type,
+    FailureActionStepExpression=FailureAction.from_string,
 )
 
 
@@ -142,6 +169,9 @@ def step_setup_log_all_requests(context: Context) -> None:
 def step_setup_stop_user_on_failure(context: Context) -> None:
     """Stop user if a request fails.
 
+    !!! attention
+        This step is deprecated and will be removed in the future, use {@pylink grizzly.steps.scenario.setup.step_setup_failed_task_default} instead.
+
     Default behavior is to continue the scenario if a request fails.
 
     Example:
@@ -151,13 +181,16 @@ def step_setup_stop_user_on_failure(context: Context) -> None:
 
     """
     grizzly = cast(GrizzlyContext, context.grizzly)
-    grizzly.scenario.failure_exception = StopUser
+    grizzly.scenario.failure_handling.update({None: StopUser})
 
 
 @given('restart scenario on failure')
 def step_setup_restart_scenario_on_failure(context: Context) -> None:
     """Restart scenario, from first task, if a request fails.
 
+    !!! attention
+        This step is deprecated and will be removed in the future, use {@pylink grizzly.steps.scenario.setup.step_setup_failed_task_default} instead.
+
     Default behavior is to continue the scenario if a request fails.
 
     Example:
@@ -167,7 +200,43 @@ def step_setup_restart_scenario_on_failure(context: Context) -> None:
 
     """
     grizzly = cast(GrizzlyContext, context.grizzly)
-    grizzly.scenario.failure_exception = RestartScenario
+    grizzly.scenario.failure_handling.update({None: RestartScenario})
+
+
+@when('a task fails with "{failure:FailureType}" {failure_action:FailureActionStepExpression}')
+def step_setup_failed_task_custom(context: Context, failure: type[Exception] | str, failure_action: FailureAction) -> None:
+    """Set behavior when specific failure occurs.
+
+    It can be either a `str` or an exception type, where the later is more specific.
+
+    Example:
+    ```gherkin
+    When a task fails with "504 gateway timeout" retry step
+    When a task fails with "RuntimeError" stop user
+    ```
+
+    """
+    grizzly = cast(GrizzlyContext, context.grizzly)
+    grizzly.scenario.failure_handling.update({failure: failure_action.exception})
+
+
+@when('a task fails {failure_action:FailureActionStepExpression}')
+def step_setup_failed_task_default(context: Context, failure_action: FailureAction) -> None:
+    """Set default behavior when a task fails.
+
+    If no default behavior is set, the scenario will continue as nothing happened.
+
+    Example:
+    ```gherkin
+    When a task fails restart scenario
+    When a task fails stop user
+    ```
+
+    """
+    assert failure_action.default_friendly, f'{failure_action.step_expression} should not be used as the default behavior, only use it for specific failures'
+
+    grizzly = cast(GrizzlyContext, context.grizzly)
+    grizzly.scenario.failure_handling.update({None: failure_action.exception})
 
 
 @then('metadata "{key}" is "{value}"')

grizzly/tasks/async_group.py

@@ -33,6 +33,7 @@
 
 import gevent
 
+from grizzly.exceptions import failure_handler
 from grizzly.types import RequestType
 from grizzly.users import AsyncRequests
 
@@ -66,7 +67,7 @@ def peek(self) -> list[GrizzlyTask]:
 
     def __call__(self) -> grizzlytask:  # noqa: C901
         @grizzlytask
-        def task(parent: GrizzlyScenario) -> Any:  # noqa: C901
+        def task(parent: GrizzlyScenario) -> Any:
             if not isinstance(parent.user, AsyncRequests):
                 message = f'{parent.user.__class__.__name__} does not inherit AsyncRequests'
                 raise NotImplementedError(message)  # pragma: no cover
@@ -136,7 +137,6 @@ def trace_green(event: str, args: tuple[gevent.Greenlet, gevent.Greenlet]) -> No
                     exception=exception,
                 )
 
-                if exception is not None and parent.user._scenario.failure_exception is not None:
-                    raise parent.user._scenario.failure_exception from exception
+                failure_handler(exception, parent.user._scenario)
 
         return task

grizzly/tasks/clients/__init__.py

@@ -29,7 +29,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast, final
 from urllib.parse import unquote, urlparse
 
-from grizzly.exceptions import StopScenario
+from grizzly.exceptions import StopScenario, failure_handler
 from grizzly.tasks import GrizzlyMetaRequestTask, grizzlytask, template
 from grizzly.testdata.utils import resolve_variable
 from grizzly.types import GrizzlyResponse, RequestDirection, RequestMethod, RequestType
@@ -304,9 +304,9 @@ def action(self, parent: GrizzlyScenario, action: Optional[str] = None, *, suppr
 
                 log_file.write_text(jsondumps(request_log, indent=2))
 
-            if exception is not None and parent.user._scenario.failure_exception is not None:
-                raise parent.user._scenario.failure_exception
-            elif exception is not None:
+            failure_handler(exception, parent.user._scenario)
+
+            if exception is not None:
                 parent.logger.warning('%s ignoring %s', self.__class__.__name__, exception)
 
 

grizzly/tasks/conditional.py

@@ -37,7 +37,7 @@
 
 from gevent import sleep as gsleep
 
-from grizzly.exceptions import RestartScenario, StopUser
+from grizzly.exceptions import RestartScenario, StopUser, failure_handler
 
 from . import GrizzlyTask, GrizzlyTaskWrapper, grizzlytask, template
 
@@ -142,8 +142,7 @@ def task(parent: GrizzlyScenario) -> Any:
                     stats = parent.user.environment.stats.get(name, 'COND')
                     stats.log_error(None)
 
-                if exception is not None and parent.user._scenario.failure_exception is not None:
-                    raise parent.user._scenario.failure_exception from exception
+                failure_handler(exception, parent.user._scenario)
 
         @task.on_start
         def on_start(parent: GrizzlyScenario) -> None:

grizzly/tasks/keystore.py

@@ -35,6 +35,8 @@
 from json import loads as jsonloads
 from typing import TYPE_CHECKING, Any, Literal, cast, get_args
 
+from grizzly.exceptions import failure_handler
+
 from . import GrizzlyTask, grizzlytask, template
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -72,7 +74,7 @@ def __init__(self, key: str, action: Action, action_context: str | Any, default_
 
     def __call__(self) -> grizzlytask:
         @grizzlytask
-        def task(parent: GrizzlyScenario) -> Any:  # noqa: PLR0912
+        def task(parent: GrizzlyScenario) -> Any:
             key = parent.user.render(self.key)
 
             try:
@@ -120,7 +122,6 @@ def task(parent: GrizzlyScenario) -> Any:  # noqa: PLR0912
                     exception=e,
                 )
 
-                if parent.user._scenario.failure_exception is not None:
-                    raise parent.user._scenario.failure_exception from e
+                failure_handler(e, parent.user._scenario)
 
         return task