Add add_test_hook() to run a hook around each function test

README.md

@@ -268,6 +268,30 @@ def test_example_3():
 
 ```
 
+### Test hooks
+
+`add_test_hook(hook)` runs `hook(env, test)` around every function test that
+takes an `env`, where `test` is the test function. The hook returns a context
+manager; failed assertions and exceptions in it fail the test. A top-level
+function can be added from a test file; add a lambda, closure or other callable
+from a module that is imported once. For example, to log how long each test
+takes:
+
+```python
+import contextlib
+import time
+from RLTest import add_test_hook
+
+@contextlib.contextmanager
+def timed(env, test):
+    start = time.monotonic()
+    yield
+    elapsed = time.monotonic() - start
+    env.debugPrint('%s took %.1f s' % (test.__name__, elapsed), force=True)
+
+add_test_hook(timed)
+```
+
 ### Shutdown grace period
 
 With the default shutdown policy (`terminateRetries=None`), RLTest retries

RLTest/__init__.py

@@ -1,12 +1,14 @@
 from RLTest.env import Env, Defaults
 from RLTest.env_spec import env_spec
 from RLTest.redis_std import StandardEnv
+from RLTest.test_hooks import add_test_hook
 from ._version import __version__
 
 __all__ = [
     'Defaults',
     'Env',
     'StandardEnv',
+    'add_test_hook',
     'env_spec',
 ]
 

RLTest/__main__.py

@@ -14,6 +14,7 @@
 import json
 from multiprocessing import Process, Queue, set_start_method
 
+from RLTest.test_hooks import run_with_hooks
 from RLTest.env import Env, TestAssertionFailure, Defaults
 from RLTest.utils import Colors, fix_modules, fix_modulesArgs, is_github_actions
 from RLTest.loader import TestLoader
@@ -746,7 +747,7 @@ def _runTest(self, test, numberOfAssertionFailed=0, prefix='', before=lambda x=N
                 return 0
 
         if env is not None:
-            fn = lambda: test.target(env)
+            fn = lambda: run_with_hooks(test.target, env)
             before_func = lambda: before(env)
             after_func = lambda: after(env)
         else:

RLTest/test_hooks.py

@@ -0,0 +1,44 @@
+"""Hooks that run around each function test that takes an ``env``."""
+import contextlib
+import inspect
+
+_hooks = []
+
+
+def add_test_hook(hook):
+    """
+    Run ``hook(env, test)`` around every function test that takes an
+    ``env``; ``test`` is the test function, so a hook can honour markers set
+    on it.
+
+    ``hook`` returns a context manager, such as a generator decorated with
+    ``contextlib.contextmanager``: it is entered before the test and exited
+    after it. Failed assertions and exceptions in the hook fail the test, as
+    the test's own would. Hooks run in the order added.
+
+    RLTest re-imports a test file before each of its tests, so a top-level
+    function added again with the same module and name replaces the copy
+    added before. Any other hook (a lambda, closure, bound method or other
+    callable) counts as the same only if it is the same object, so add those
+    from a module that is imported once.
+    """
+    key = _key(hook)
+    for i, added in enumerate(_hooks):
+        if added is hook or _key(added) == key:
+            _hooks[i] = hook
+            return
+    _hooks.append(hook)
+
+
+def _key(hook):
+    name = getattr(hook, '__qualname__', '')
+    if inspect.isfunction(hook) and name.isidentifier():
+        return (hook.__module__, name)
+    return (id(hook),)
+
+
+def run_with_hooks(target, env):
+    with contextlib.ExitStack() as stack:
+        for hook in list(_hooks):
+            stack.enter_context(hook(env, target))
+        return target(env)

tests/flow/test_hooks_flow.py

@@ -0,0 +1,23 @@
+import contextlib
+
+from RLTest import add_test_hook
+
+
+# Added at import, as a test file would. RLTest re-imports this file before
+# each of its tests, so a hook added more than once would run more than once.
+@contextlib.contextmanager
+def count_calls(env, test):
+    calls = env.__dict__.setdefault('_hook_calls', {})
+    calls[test.__name__] = calls.get(test.__name__, 0) + 1
+    yield
+
+
+add_test_hook(count_calls)
+
+
+def test_hook_runs_once_per_test(env):
+    env.assertEqual(env._hook_calls.get('test_hook_runs_once_per_test'), 1)
+
+
+def test_hook_runs_once_per_later_test(env):
+    env.assertEqual(env._hook_calls.get('test_hook_runs_once_per_later_test'), 1)

tests/unit/test_test_hooks.py

@@ -0,0 +1,130 @@
+"""Unit tests for add_test_hook."""
+import contextlib
+import types
+
+import pytest
+
+from RLTest import add_test_hook
+from RLTest import test_hooks
+from RLTest.test_hooks import run_with_hooks
+
+
+@pytest.fixture(autouse=True)
+def no_hooks():
+    saved = list(test_hooks._hooks)
+    del test_hooks._hooks[:]
+    yield
+    test_hooks._hooks[:] = saved
+
+
+def recorder(log, name):
+    @contextlib.contextmanager
+    def hook(env, test):
+        log.append((name, 'enter', env))
+        try:
+            yield
+        except BaseException as error:
+            log.append((name, 'exit', type(error).__name__))
+            raise
+        log.append((name, 'exit', env))
+    return hook
+
+
+TEST_FILE = '''
+import contextlib
+
+@contextlib.contextmanager
+def hook(env, test):
+    LOG.append(TAG)
+    yield
+'''
+
+
+def import_test_file(log, tag):
+    """Run TEST_FILE as RLTest's loader does, as a fresh module each time."""
+    module = types.ModuleType('test_file')
+    module.LOG, module.TAG = log, tag
+    exec(compile(TEST_FILE, 'test_file.py', 'exec'), module.__dict__)
+    return module.hook
+
+
+def test_no_hooks_runs_the_test():
+    assert run_with_hooks(lambda env: env * 2, 21) == 42
+
+
+def test_hooks_wrap_the_test_in_order():
+    log = []
+    add_test_hook(recorder(log, 'a'))
+    add_test_hook(recorder(log, 'b'))
+    run_with_hooks(lambda env: log.append(('test', env)), 'env')
+    assert log == [('a', 'enter', 'env'), ('b', 'enter', 'env'), ('test', 'env'),
+                   ('b', 'exit', 'env'), ('a', 'exit', 'env')]
+
+
+def test_adding_a_hook_twice_runs_it_once():
+    log = []
+    hook = recorder(log, 'a')
+    add_test_hook(hook)
+    add_test_hook(hook)
+    run_with_hooks(lambda env: None, 'env')
+    assert log == [('a', 'enter', 'env'), ('a', 'exit', 'env')]
+
+
+def test_a_hook_from_a_reimported_test_file_replaces_the_old_one():
+    log = []
+    add_test_hook(import_test_file(log, 'first import'))
+    add_test_hook(recorder(log, 'other'))
+    add_test_hook(import_test_file(log, 'second import'))
+    run_with_hooks(lambda env: None, 'env')
+    # One copy, the newer, still first.
+    assert log == ['second import', ('other', 'enter', 'env'), ('other', 'exit', 'env')]
+
+
+def test_distinct_hooks_that_share_a_name_all_run():
+    log = []
+    add_test_hook(recorder(log, 'a'))  # closures from one factory
+    add_test_hook(recorder(log, 'b'))
+    add_test_hook(lambda env, test: recorder(log, 'c')(env, test))
+    add_test_hook(lambda env, test: recorder(log, 'd')(env, test))
+    run_with_hooks(lambda env: None, 'env')
+    assert [name for name, step, _ in log if step == 'enter'] == ['a', 'b', 'c', 'd']
+
+
+def test_a_failing_test_still_exits_the_hooks():
+    log = []
+    add_test_hook(recorder(log, 'a'))
+
+    def failing(env):
+        raise ValueError('boom')
+
+    with pytest.raises(ValueError, match='boom'):
+        run_with_hooks(failing, 'env')
+    assert log == [('a', 'enter', 'env'), ('a', 'exit', 'ValueError')]
+
+
+def test_a_hook_error_fails_the_test():
+    @contextlib.contextmanager
+    def check(env, test):
+        yield
+        raise AssertionError('hook check failed')
+
+    add_test_hook(check)
+    with pytest.raises(AssertionError, match='hook check failed'):
+        run_with_hooks(lambda env: None, 'env')
+
+
+def test_a_hook_gets_the_test_function():
+    seen = []
+
+    @contextlib.contextmanager
+    def hook(env, test):
+        seen.append(getattr(test, 'marker', None))
+        yield
+
+    def marked(env):
+        pass
+    marked.marker = 'skip-check'
+
+    add_test_hook(hook)
+    run_with_hooks(marked, 'env')
+    assert seen == ['skip-check']