Repository: taverntesting/tavern

Issue:
Add support for before/after test stage actions
follow up to this comment: https://github.com/taverntesting/tavern/issues/41#issuecomment-370462578

Currently, the support for pre- and post-stage actions is rather limited. The only ones we support right now are delays and clearing cookies. I suggest adding the possibility to call arbitrary functions here. The format could be similar to the `$ext` function calls currently allowed for validation.

One use case we have for this would be to use [freezegun](https://github.com/spulec/freezegun) to simulate changes in time for a single request (e.g. for testing session expiry).

```yaml
stages:
  - name: test-datetime-endpoint

    before:
      - function: my.package.utils.freeze
        extra_kwargs:
          time: "2018-05-31 12:00:00"

    after:
      - function: my.package.utils.unfreeze

    request:
      url: "{host}/get-time"
      method: GET

    response:
      status_code: 200
      body:
        current_time: "2018-05-31 12:00:00"
```

```python
# my.package.utils

def freeze(time, context):
    t = context['freezegun'] = freezegun.freeze_time(time)
    t.start()


def unfreeze(context):
    t = context.pop('freezegun')
    t.stop()

```

Note the `context` parameter I have added here which allows passing a context object (might just be a dictionary) to the function. This allows a pair of setup- and teardown methods to keep track of created objects. The context may also contain additional info, for example the session is required for clearing the cookies.

The context variable should be injected dynamically, which is controlled either through an attribute in the spec or through a decorator:

```yaml
      - function: my.package.utils.freeze
        with_context: true
        extra_kwargs:
          time: "2018-05-31 12:00:00"
```

```python
# sets freeze._context = True
@tavern.utils.inject_context
def freeze(time, context):
    t = context['freezegun'] = freezegun.freeze_time(time)
    t.start()
```

This allows us to still use functions without context which is useful when using builtins (e.g. `time.sleep` as a replacement for `delay_after` and `delay_before`).

What are your thoughts on this?



Discussion:

Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 2e20913..f42a133 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -47,6 +47,10 @@ Run every so often to update the pre-commit hooks
     black tavern/ tests/
     ruff --fix tavern/ tests/
 
+### Fix yaml formatting issues
+
+    pre-commit run --all-files
+
 ## Creating a new release
 
 1. Setup `~/.pypirc`
diff --git a/docs/source/basics.md b/docs/source/basics.md
index ac243d9..e0c3793 100644
--- a/docs/source/basics.md
+++ b/docs/source/basics.md
@@ -2381,6 +2381,79 @@ def pytest_tavern_beta_before_every_request(request_args):
     logging.info("Making request: %s", request_args)
 ```
 
+## Tinctures
+
+Another way of running functions at certain times is to use the 'tinctures' functionality:
+
+```python
+# package/helpers.py
+
+import logging
+import time
+
+logger = logging.getLogger(__name__)
+
+
+def time_request(stage):
+    t0 = time.time()
+    yield
+    t1 = time.time()
+    logger.info("Request for stage %s took %s", stage, t1 - t0)
+
+
+def print_response(_, extra_print="affa"):
+    logger.info("STARTING:")
+    (expected, response) = yield
+    logger.info("Response is %s (%s)", response, extra_print)
+```
+
+```yaml
+---
+test_name: Test tincture
+
+tinctures:
+  - function: package.helpers:time_request
+
+stages:
+  - name: Make a request
+    tinctures:
+      - function: package.helpers:print_response
+        extra_kwargs:
+          extra_print: "blooble"
+    request:
+      url: "{host}/echo"
+      method: POST
+      json:
+        value: "one"
+
+  - name: Make another request
+    request:
+      url: "{host}/echo"
+      method: POST
+      json:
+        value: "two"
+```
+
+Tinctures can be specified on a per-stage level or a per-test level. When specified on the test level, the tincture is
+run for every stage in the test. In the above example, the `time_request` function will be run for both stages, but
+the 'print_response' function will only be run for the first stage.
+
+Tinctures are _similar_ to fixtures but are more similar to [external functions](#calling-external-functions). Tincture
+functions do not need to be annotated with a function like Pytest fixtures, and are referred to in the same
+way (`path.to.package:function`), and have arguments passed to them in the same way (`extra_kwargs`, `extra_args`) as
+external functions.
+
+The first argument to a tincture is always a dictionary of the stage to be run.
+
+If a tincture has a `yield` in the middle of it, during the `yield` the stage itself will be run. If a return value is
+expected from the `yield` (eg `(expected, response) = yield` in the example above) then the _expected_ return values and
+the response object from the stage will be returned. This allows a tincture to introspect the response, and compare it
+against the expected, the same as the `pytest_tavern_beta_after_every_response` [hook](#after-every-response). This
+response object will be different for MQTT and HTTP tests!
+
+If you need to run something before _every_ stage or after _every_ response in your test suite, look at using
+the [hooks](#hooks) instead.
+
 ## Finalising stages
 
 If you need a stage to run after a test runs, whether it passes or fails (for example, to log out of a service or
@@ -2396,7 +2469,7 @@ stages:
 
   - name: stage 2
     ...
-    
+
   - name: stage 3
     ...
 
diff --git a/requirements.txt b/requirements.txt
index 3ecae75..d643f72 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -316,6 +316,10 @@ docutils==0.18.1 \
     #   sphinx
     #   sphinx-rtd-theme
     #   tavern (pyproject.toml)
+exceptiongroup==1.1.3 \
+    --hash=sha256:097acd85d473d75af5bb98e41b61ff7fe35efe6675e4f9370ec6ec5126d160e9 \
+    --hash=sha256:343280667a4585d195ca1cf9cef84a4e178c4b6cf2274caef9859782b567d5e3
+    # via pytest
 execnet==2.0.2 \
     --hash=sha256:88256416ae766bc9e8895c76a87928c0012183da3cc4fc18016e6f050e025f41 \
     --hash=sha256:cc59bc4423742fd71ad227122eb0dd44db51efb3dc4095b45ac9a08c770096af
@@ -930,6 +934,18 @@ stevedore==4.1.1 \
     --hash=sha256:7f8aeb6e3f90f96832c301bff21a7eb5eefbe894c88c506483d355565d88cc1a \
     --hash=sha256:aa6436565c069b2946fe4ebff07f5041e0c8bf18c7376dd29edf80cf7d524e4e
     # via tavern (pyproject.toml)
+tomli==2.0.1 \
+    --hash=sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc \
+    --hash=sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f
+    # via
+    #   build
+    #   coverage
+    #   mypy
+    #   pip-tools
+    #   pyproject-api
+    #   pyproject-hooks
+    #   pytest
+    #   tox
 tomli-w==1.0.0 \
     --hash=sha256:9f2a07e8be30a0729e533ec968016807069991ae2fd921a78d42f429ae5f4463 \
     --hash=sha256:f463434305e0336248cac9c2dc8076b707d8a12d019dd349f5c1e382dd1ae1b9
diff --git a/tavern/_core/dict_util.py b/tavern/_core/dict_util.py
index 3baea4e..fd0ab1c 100644
--- a/tavern/_core/dict_util.py
+++ b/tavern/_core/dict_util.py
@@ -111,7 +111,7 @@ def format_keys(
             in broken output, only use for debugging purposes.
 
     Returns:
-        str,int,list,dict: recursively formatted values
+        recursively formatted values
     """
     formatted = val
 
@@ -131,7 +131,7 @@ def format_keys(
         for key in val:
             formatted[key] = format_keys_(val[key], box_vars)
     elif isinstance(val, (list, tuple)):
-        formatted = [format_keys_(item, box_vars) for item in val]
+        formatted = [format_keys_(item, box_vars) for item in val]  # type: ignore
     elif isinstance(formatted, FormattedString):
         logger.debug("Already formatted %s, not double-formatting", formatted)
     elif isinstance(val, str):
@@ -142,7 +142,7 @@ def format_keys(
                 raise
 
         if no_double_format:
-            formatted = FormattedString(formatted)
+            formatted = FormattedString(formatted)  # type: ignore
     elif isinstance(val, TypeConvertToken):
         logger.debug("Got type convert token '%s'", val)
         if isinstance(val, ForceIncludeToken):
diff --git a/tavern/_core/run.py b/tavern/_core/run.py
index 2b42e5c..50a3535 100644
--- a/tavern/_core/run.py
+++ b/tavern/_core/run.py
@@ -25,6 +25,7 @@ from .pytest.config import TestConfig
 from .report import attach_stage_content, wrap_step
 from .strtobool import strtobool
 from .testhelpers import delay, retry
+from .tincture import Tinctures, get_stage_tinctures
 
 logger = logging.getLogger(__name__)
 
@@ -276,6 +277,8 @@ class _TestRunner:
     test_spec: Mapping
 
     def run_stage(self, idx: int, stage, *, is_final: bool = False):
+        tinctures = get_stage_tinctures(stage, self.test_spec)
+
         stage_config = self.test_block_config.with_strictness(
             self.default_global_strictness
         )
@@ -284,7 +287,9 @@ class _TestRunner:
         )
         # Wrap run_stage with retry helper
         run_stage_with_retries = retry(stage, stage_config)(self.wrapped_run_stage)
-        partial = functools.partial(run_stage_with_retries, stage, stage_config)
+        partial = functools.partial(
+            run_stage_with_retries, stage, stage_config, tinctures
+        )
         allure_name = "Stage {}: {}".format(
             idx, format_keys(stage["name"], stage_config.variables)
         )
@@ -298,12 +303,15 @@ class _TestRunner:
             e.is_final = is_final
             raise
 
-    def wrapped_run_stage(self, stage: dict, stage_config: TestConfig):
+    def wrapped_run_stage(
+        self, stage: dict, stage_config: TestConfig, tinctures: Tinctures
+    ):
         """Run one stage from the test
 
         Args:
             stage: specification of stage to be run
             stage_config: available variables for test
+            tinctures: tinctures for this stage/test
         """
         stage = copy.deepcopy(stage)
         name = stage["name"]
@@ -329,8 +337,12 @@ class _TestRunner:
 
         verifiers = get_verifiers(stage, stage_config, self.sessions, expected)
 
+        tinctures.start_tinctures(stage)
+
         response = r.run()
 
+        tinctures.end_tinctures(expected, response)
+
         for response_type, response_verifiers in verifiers.items():
             logger.debug("Running verifiers for %s", response_type)
             for v in response_verifiers:
diff --git a/tavern/_core/tincture.py b/tavern/_core/tincture.py
new file mode 100644
index 0000000..ae33dbf
--- /dev/null
+++ b/tavern/_core/tincture.py
@@ -0,0 +1,81 @@
+import collections.abc
+import inspect
+import logging
+from typing import Any, List
+
+from tavern._core import exceptions
+from tavern._core.extfunctions import get_wrapped_response_function
+
+logger = logging.getLogger(__name__)
+
+
+class Tinctures:
+    def __init__(self, tinctures: List[Any]):
+        self._tinctures = tinctures
+        self._needs_response: List[Any] = []
+
+    def start_tinctures(self, stage: collections.abc.Mapping):
+        results = [t(stage) for t in self._tinctures]
+        self._needs_response = []
+
+        for r in results:
+            if inspect.isgenerator(r):
+                # Store generator and start it
+                self._needs_response.append(r)
+                next(r)
+
+    def end_tinctures(self, expected: collections.abc.Mapping, response) -> None:
+        """
+        Send the response object to any tinctures that want it
+
+        Args:
+            response: The response from 'run' for the stage
+        """
+        if self._needs_response is None:
+            raise RuntimeError(
+                "should not be called before accumulating tinctures which need a response"
+            )
+
+        for n in self._needs_response:
+            try:
+                n.send((expected, response))
+            except StopIteration:
+                pass
+            else:
+                raise RuntimeError("Tincture had more than one yield")
+
+
+def get_stage_tinctures(
+    stage: collections.abc.Mapping, test_spec: collections.abc.Mapping
+) -> Tinctures:
+    """Get tinctures for stage
+
+    Args:
+        stage: Stage
+        test_spec: Whole test spec
+    """
+    stage_tinctures = []
+
+    def add_tinctures_from_block(maybe_tinctures, blockname: str):
+        logger.debug("Trying to add tinctures from %s", blockname)
+
+        def inner_yield():
+            if maybe_tinctures is not None:
+                if isinstance(maybe_tinctures, list):
+                    for vf in maybe_tinctures:
+                        yield get_wrapped_response_function(vf)
+                elif isinstance(maybe_tinctures, dict):
+                    yield get_wrapped_response_function(maybe_tinctures)
+                elif maybe_tinctures is not None:
+                    raise exceptions.BadSchemaError(
+                        "Badly formatted 'tinctures' block in {}".format(blockname)
+                    )
+
+        stage_tinctures.extend(inner_yield())
+
+    add_tinctures_from_block(test_spec.get("tinctures"), "test")
+    add_tinctures_from_block(stage.get("tinctures"), "stage")
+
+    logger.debug("%d tinctures for stage %s", len(stage_tinctures), stage["name"])
+
+    return Tinctures(stage_tinctures)


Test patch:
diff --git a/tavern/_core/schema/tests.jsonschema.yaml b/tavern/_core/schema/tests.jsonschema.yaml
index 84a0899..473856e 100644
--- a/tavern/_core/schema/tests.jsonschema.yaml
+++ b/tavern/_core/schema/tests.jsonschema.yaml
@@ -322,6 +322,12 @@ definitions:
       - name
 
     properties:
+      tinctures:
+        type: array
+        description: Tinctures for stage
+        items:
+          $ref: "#/definitions/verify_block"
+
       id:
         type: string
         description: ID of stage for use in stage references
@@ -415,6 +421,12 @@ properties:
                 - key
                 - vals
 
+  tinctures:
+    type: array
+    description: Tinctures for whole test
+    items:
+      $ref: "#/definitions/verify_block"
+
   strict:
     $ref: "#/definitions/strict_block"
 
diff --git a/tavern/_core/schema/tests.schema.yaml b/tavern/_core/schema/tests.schema.yaml
index 4cc2639..ced99f6 100644
--- a/tavern/_core/schema/tests.schema.yaml
+++ b/tavern/_core/schema/tests.schema.yaml
@@ -28,6 +28,10 @@ schema;stage:
   required: true
   func: verify_oneof_id_name
   mapping:
+    tinctures:
+      func: validate_extensions
+      type: any
+
     id:
       type: str
       required: false
diff --git a/tavern/_core/testhelpers.py b/tavern/_core/testhelpers.py
index 095c347..42247f6 100644
--- a/tavern/_core/testhelpers.py
+++ b/tavern/_core/testhelpers.py
@@ -10,13 +10,13 @@ from tavern._core.pytest.config import TestConfig
 logger = logging.getLogger(__name__)
 
 
-def delay(stage, when, variables) -> None:
+def delay(stage: Mapping, when: str, variables: Mapping) -> None:
     """Look for delay_before/delay_after and sleep
 
     Args:
-        stage (dict): test stage
-        when (str): 'before' or 'after'
-        variables (dict): Variables to format with
+        stage: test stage
+        when: 'before' or 'after'
+        variables: Variables to format with
     """
 
     try:
diff --git a/tests/integration/ext_functions.py b/tests/integration/ext_functions.py
index fbb23ab..cf33102 100644
--- a/tests/integration/ext_functions.py
+++ b/tests/integration/ext_functions.py
@@ -1,3 +1,6 @@
+import time
+
+
 def return_hello():
     return {"hello": "there"}
 
@@ -12,3 +15,13 @@ def return_list_vals():
 
 def gen_echo_url(host):
     return "{0}/echo".format(host)
+
+
+def time_request(_):
+    time.time()
+    yield
+    time.time()
+
+
+def print_response(_, extra_print="affa"):
+    (_, r) = yield
diff --git a/tests/integration/test_tincture.tavern.yaml b/tests/integration/test_tincture.tavern.yaml
new file mode 100644
index 0000000..06bd178
--- /dev/null
+++ b/tests/integration/test_tincture.tavern.yaml
@@ -0,0 +1,60 @@
+---
+test_name: Test tincture + fixtures
+
+includes:
+  - !include common.yaml
+
+tinctures:
+  - function: ext_functions:time_request
+  - function: ext_functions:print_response
+    extra_kwargs:
+      extra_print: "blooble"
+
+marks:
+  - usefixtures:
+      - yielder
+      - str_fixture
+      - yield_str_fixture
+
+stages:
+  - name: do something
+    tinctures:
+      - function: ext_functions:time_request
+      - function: ext_functions:print_response
+        extra_kwargs:
+          extra_print: "blooble"
+    request:
+      url: "{host}/echo"
+      method: POST
+      json:
+        value: "{str_fixture}"
+    response:
+      status_code: 200
+      json:
+        value: "{yield_str_fixture}"
+
+---
+test_name: Test tincture extra kwargs fails
+
+includes:
+  - !include common.yaml
+
+tinctures:
+  - function: ext_functions:print_response
+    extra_kwargs:
+      extra_print: "blooble"
+      something: else
+
+_xfail: run
+
+stages:
+  - name: do something
+    request:
+      url: "{host}/echo"
+      method: POST
+      json:
+        value: "{str_fixture}"
+    response:
+      status_code: 200
+      json:
+        value: "{yield_str_fixture}"
diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py
index 1751538..e930d81 100644
--- a/tests/unit/test_core.py
+++ b/tests/unit/test_core.py
@@ -587,6 +587,52 @@ class TestFinally:
         assert pmock.mock_calls[1].kwargs.items() >= finally_request["request"].items()
 
 
+class TestTinctures:
+    @pytest.mark.parametrize(
+        "tinctures",
+        (
+            {"function": "abc"},
+            [{"function": "abc"}],
+            [{"function": "abc"}, {"function": "def"}],
+        ),
+    )
+    @pytest.mark.parametrize(
+        "at_stage_level",
+        (
+            True,
+            False,
+        ),
+    )
+    def test_tinctures(
+        self,
+        fulltest,
+        mockargs,
+        includes,
+        tinctures,
+        at_stage_level,
+    ):
+        if at_stage_level:
+            fulltest["tinctures"] = tinctures
+        else:
+            fulltest["stages"][0]["tinctures"] = tinctures
+
+        mock_response = Mock(**mockargs)
+
+        tincture_func_mock = Mock()
+
+        with patch(
+            "tavern._plugins.rest.request.requests.Session.request",
+            return_value=mock_response,
+        ) as pmock, patch(
+            "tavern._core.tincture.get_wrapped_response_function",
+            return_value=tincture_func_mock,
+        ):
+            run_test("heif", fulltest, includes)
+
+        assert pmock.call_count == 1
+        assert tincture_func_mock.call_count == len(tinctures)
+
+
 def test_copy_config(pytestconfig):
     cfg_1 = load_global_cfg(pytestconfig)
 
diff --git a/tests/unit/test_tinctures.py b/tests/unit/test_tinctures.py
new file mode 100644
index 0000000..98768a7
--- /dev/null
+++ b/tests/unit/test_tinctures.py
@@ -0,0 +1,113 @@
+from unittest.mock import patch
+
+import pytest
+
+from tavern._core.tincture import Tinctures, get_stage_tinctures
+
+
+@pytest.fixture(name="example")
+def example():
+    spec = {
+        "test_name": "A test with a single stage",
+        "stages": [
+            {
+                "name": "step 1",
+                "request": {"url": "http://www.google.com", "method": "GET"},
+                "response": {
+                    "status_code": 200,
+                    "json": {"key": "value"},
+                    "headers": {"content-type": "application/json"},
+                },
+            }
+        ],
+    }
+
+    return spec
+
+
+def test_empty():
+    t = Tinctures([])
+
+    t.start_tinctures({})
+    t.end_tinctures({}, None)
+
+
+@pytest.mark.parametrize(
+    "tinctures",
+    (
+        {"function": "abc"},
+        [{"function": "abc"}],
+        [{"function": "abc"}, {"function": "def"}],
+    ),
+)
+class TestTinctures:
+    class TestTinctureBasic:
+        def test_stage_tinctures_normal(self, example, tinctures):
+            stage = example["stages"][0]
+            stage["tinctures"] = tinctures
+
+            with patch(
+                "tavern._core.tincture.get_wrapped_response_function",
+                return_value=lambda _: None,
+            ) as call_mock:
+                t = get_stage_tinctures(stage, example)
+
+            t.start_tinctures(stage)
+            t.end_tinctures(stage, None)
+
+            assert call_mock.call_count == len(tinctures)
+
+        def test_test_tinctures_normal(self, example, tinctures):
+            stage = example["stages"][0]
+            example["tinctures"] = tinctures
+
+            with patch(
+                "tavern._core.tincture.get_wrapped_response_function",
+                return_value=lambda _: None,
+            ) as call_mock:
+                t = get_stage_tinctures(stage, example)
+
+            t.start_tinctures(stage)
+            t.end_tinctures(stage, None)
+
+            assert call_mock.call_count == len(tinctures)
+
+    class TestTinctureYields:
+        @staticmethod
+        def does_yield(stage):
+            assert stage["name"] == "step 1"
+
+            (expected, response) = yield
+
+            assert expected == stage["response"]
+            assert response is None
+
+        def test_stage_tinctures_normal(self, example, tinctures):
+            stage = example["stages"][0]
+            stage["tinctures"] = tinctures
+
+            with patch(
+                "tavern._core.tincture.get_wrapped_response_function",
+                return_value=self.does_yield,
+            ) as call_mock:
+                t = get_stage_tinctures(stage, example)
+
+            t.start_tinctures(stage)
+            t.end_tinctures(stage["response"], None)
+
+            assert call_mock.call_count == len(tinctures)
+
+        def test_test_tinctures_normal(self, example, tinctures):
+            stage = example["stages"][0]
+            example["tinctures"] = tinctures
+
+            with patch(
+                "tavern._core.tincture.get_wrapped_response_function",
+                return_value=self.does_yield,
+            ) as call_mock:
+                t = get_stage_tinctures(stage, example)
+
+            t.start_tinctures(stage)
+            t.end_tinctures(stage["response"], None)
+
+            assert call_mock.call_count == len(tinctures)