Try supporting yield functions for setup/teardown. (#1)

* Support yield functions for setup/teardown.

Co-authored-by: Blake Naccarato <[email protected]>

.coveragerc

@@ -4,4 +4,5 @@ branch = True
 [paths]
 source =
    src
-   .tox/*/site-packages
+   .tox/*/lib/*/site-packages
+

README.md

@@ -81,6 +81,52 @@ param teardown
 * If an exception occurs in setup, the test will report Error and not run. The teardown will also not run.
 * If an exception occurs in teardown, the LAST parametrized test case to run results in BOTH PASS and Error. This is weird, but consistent with pytest fixtures.
 
+
+## You can combine setup and teardown in one function
+
+You can provide a function separated by a `yield` to put both setup and teardown in one function.
+
+However, there's a trick to doing this:
+
+* Either, pass `None` as the teardown.
+* Or use `with_args`, as in `@pytest.mark.param_scope.with_args(my_func)`
+
+Here's a combo setup/teardown function:
+
+```python
+def setup_and_teardown():
+    print('\nsetup')
+    yield 42
+    print('\nteardown')
+
+```
+
+Calling it with `None` for teardown:
+
+```python
+import pytest
+
+@pytest.mark.param_scope(setup_and_teardown, None)
+@pytest.mark.parametrize('x', ['a', 'b', 'c'])
+def test_yield(x, param_scope):
+    assert param_scope == 42
+
+```
+
+Or using `with_args`:
+
+```python
+@pytest.mark.param_scope.with_args(setup_and_teardown)
+@pytest.mark.parametrize('x', ['a', 'b', 'c'])
+def test_just_one_func(x, param_scope):
+    assert param_scope == 42
+
+```
+
+Both of these examples are in `examples/test_yield.py`.
+
+
+
 ## More examples
 
 Please see `examples` directory in the repo.

examples/test_error.py

@@ -36,4 +36,3 @@ def test_error_during_teardown(x):
     """
     ...
 
-

examples/test_marker_bad_params.py

@@ -0,0 +1,13 @@
+import pytest
+
+
+def foo():
+    ...
+
+@pytest.mark.param_scope(foo)
+def test_one_params_to_marker():
+    """
+    This also blows up, with_args required.
+    You gotta use `@pytest.mark.param_scope.with_args(foo)`
+    """
+    ...

examples/test_yield.py

@@ -0,0 +1,49 @@
+import pytest
+
+
+def setup_and_teardown():
+    print('\nsetup')
+    yield 42
+    print('\nteardown')
+
+
+@pytest.mark.param_scope(setup_and_teardown, None)
+@pytest.mark.parametrize('x', ['a', 'b', 'c'])
+def test_yield(x, param_scope):
+    assert param_scope == 42
+
+
+
+def separate_teardown():
+    print('separate teardown')
+
+
+@pytest.mark.param_scope(setup_and_teardown, separate_teardown)
+@pytest.mark.parametrize('x', ['a', 'b', 'c'])
+def test_two_teardowns(x, param_scope):
+    """
+    For now, we'll allow this odd use model.
+    Weird, but really, why not?
+    """
+    assert param_scope == 42
+
+
+@pytest.mark.param_scope.with_args(setup_and_teardown)
+@pytest.mark.parametrize('x', ['a', 'b', 'c'])
+def test_just_one_func(x, param_scope):
+    """
+    It's not pretty, but if you want to just pass in one,
+    you gotta use "with_args".
+    See "Passing a callable to custom markers" in pytest docs
+    - https://docs.pytest.org/en/stable/example/markers.html#passing-a-callable-to-custom-markers
+    """
+    assert param_scope == 42
+
+
+@pytest.mark.param_scope
+@pytest.mark.parametrize('x', ['a', 'b', 'c'])
+def test_no_param_scope_args(x):
+    """
+    No point in this, but it doesn't blow up
+    """
+    ...

src/pytest_param_scope/plugin.py

@@ -1,4 +1,6 @@
 from __future__ import annotations
+import types
+from typing import Generator
 import pytest
 from dataclasses import dataclass
 from typing import Callable, Any
@@ -14,6 +16,7 @@ def pytest_configure(config):
 class ParamScopeData():
     test_name: str | None = None
     teardown_func: Callable | None = None
+    teardown_gen: Generator[Any, None, None] | None = None
     ready_for_teardown: bool = False
     setup_value: Any = None
     exception: Exception | None = None
@@ -39,12 +42,24 @@ def param_scope(request):
 
         m = request.node.get_closest_marker("param_scope")
         if m:
-            setup_func = m.args[0]
-            __data.teardown_func = m.args[1]
+            if len(m.args) >= 1:
+                setup_func = m.args[0]
+            if len(m.args) >= 2:
+                __data.teardown_func = m.args[1]
 
         if setup_func:
             try:
-                __data.setup_value = setup_func()
+                # setup could be a func, or could be a generator
+                setup_value = setup_func()
+                if isinstance(setup_value, types.GeneratorType):
+                    # if generator, call next() once for setup section
+                    new_value = next(setup_value)
+                    __data.setup_value = new_value
+                    # and save it for teardown
+                    __data.teardown_gen = setup_value
+                else:
+                    # otherwise, just save the value
+                    __data.setup_value = setup_value
             except Exception as e:
                 __data.exception = e
                 raise e
@@ -58,8 +73,20 @@ def param_scope(request):
     yield __data.setup_value
 
     if __data.ready_for_teardown:
+        teardown_gen = __data.teardown_gen
         teardown_func = __data.teardown_func
+
         __data = ParamScopeData()  # reset for next one
+
+        if teardown_gen:
+            try:
+                next(teardown_gen)
+            except StopIteration:
+                pass # this is expected
+
+
+        # should we disallow both a teardown from gen and a teardown?
+        # for now, I'll allow it, as it's a bit messy to disallow it
         if teardown_func:
             teardown_func()
 

tests/test_error.py

@@ -27,3 +27,20 @@ def test_error_during_teardown(pytester):
     result = pytester.runpytest('test_error.py::test_error_during_teardown', '-v', '-s')
     result.assert_outcomes(passed=3, errors=1)
 
+def test_error_marker_bad_params(pytester):
+    """
+    Markers that accept functions have to accept 2 or more.
+
+    - all tests to pass
+    - last test to error
+    - yes, this is normal-ish for pytest with parametrized errors.
+    """
+    pytester.copy_example("examples/test_marker_bad_params.py")
+    result = pytester.runpytest('-v', '-s')
+    result.assert_outcomes(errors=1)
+    result.stdout.re_match_lines(
+        [
+            ".*Interrupted: 1 error during collection.*"
+        ]
+    )
+

tests/test_yield.py

@@ -0,0 +1,56 @@
+
+def test_yield(pytester):
+    pytester.copy_example("examples/test_yield.py")
+    result = pytester.runpytest('test_yield.py::test_yield', '-v', '-s')
+    result.assert_outcomes(passed=3)
+    result.stdout.re_match_lines(
+        [
+            ".*test_yield.a.*",
+            "setup",
+            ".*test_yield.b.*",
+            ".*test_yield.c.*",
+            "teardown",
+        ]
+    )
+
+def test_two_teardowns(pytester):
+    pytester.copy_example("examples/test_yield.py")
+    result = pytester.runpytest('test_yield.py::test_two_teardowns', '-v', '-s')
+    result.assert_outcomes(passed=3)
+    result.stdout.re_match_lines(
+        [
+            ".*test_two_teardowns.a.*",
+            "setup",
+            ".*test_two_teardowns.b.*",
+            ".*test_two_teardowns.c.*",
+            "teardown",
+            "separate teardown",
+        ]
+    )
+
+def test_one_param(pytester):
+    pytester.copy_example("examples/test_yield.py")
+    result = pytester.runpytest('test_yield.py::test_just_one_func', '-v', '-s')
+    result.assert_outcomes(passed=3)
+    result.stdout.re_match_lines(
+        [
+            ".*test_just_one_func.a.*",
+            "setup",
+            ".*test_just_one_func.b.*",
+            ".*test_just_one_func.c.*",
+            "teardown",
+        ]
+    )
+
+
+def test_no_params(pytester):
+    pytester.copy_example("examples/test_yield.py")
+    result = pytester.runpytest('test_yield.py::test_no_param_scope_args', '-v', '-s')
+    result.assert_outcomes(passed=3)
+    result.stdout.re_match_lines(
+        [
+            ".*test_no_param_scope_args.a.*",
+            ".*test_no_param_scope_args.b.*",
+            ".*test_no_param_scope_args.c.*",
+        ]
+    )