From 5a4afe9d78fa2d09915ecf8e1b135e44d3c50f71 Mon Sep 17 00:00:00 2001
From: Brian Okken <1568356+okken@users.noreply.github.com>
Date: Mon, 16 Oct 2023 08:07:01 -0700
Subject: [PATCH 1/6] Try supporting yield functions for setup/teardown.

---
 examples/test_class.py           | 14 ++++++++++++++
 examples/test_yield.py           | 13 +++++++++++++
 src/pytest_param_scope/plugin.py | 27 +++++++++++++++++++++++++--
 tests/test_error.py              |  7 +++++++
 tests/test_yield.py              | 14 ++++++++++++++
 5 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 examples/test_class.py
 create mode 100644 examples/test_yield.py
 create mode 100644 tests/test_yield.py

diff --git a/examples/test_class.py b/examples/test_class.py
new file mode 100644
index 0000000..8e7e075
--- /dev/null
+++ b/examples/test_class.py
@@ -0,0 +1,14 @@
+import pytest
+
+
+@pytest.fixture(scope='class')
+def a_fixture():
+    print('\nfixture setup')
+    yield 'some data from setup'
+    print('\nfixture teardown')
+
+
+class TestClass():
+    @pytest.mark.parametrize('x', ['a', 'b', 'c'])
+    def test_param(self, x, a_fixture):
+        assert a_fixture == 'some data from setup'
\ No newline at end of file
diff --git a/examples/test_yield.py b/examples/test_yield.py
new file mode 100644
index 0000000..50d9678
--- /dev/null
+++ b/examples/test_yield.py
@@ -0,0 +1,13 @@
+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
diff --git a/src/pytest_param_scope/plugin.py b/src/pytest_param_scope/plugin.py
index 288bb57..e0c4710 100644
--- a/src/pytest_param_scope/plugin.py
+++ b/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
@@ -44,9 +47,19 @@ def param_scope(request):
 
         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 = 3
+                __data.exception = e
                 raise e
     else:
         if __data.test_name == test_name:
@@ -58,8 +71,18 @@ 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?
         if teardown_func:
             teardown_func()
 
diff --git a/tests/test_error.py b/tests/test_error.py
index 8a3d21b..ff4b3fc 100644
--- a/tests/test_error.py
+++ b/tests/test_error.py
@@ -8,6 +8,13 @@ def test_error_during_setup(pytester):
     result = pytester.runpytest('test_error.py::test_error_during_setup', '-v', '-s')
     result.assert_outcomes(errors=3)
     result.stdout.no_fnmatch_line("param teardown")
+    result.stdout.re_match_lines(
+        [
+        "ERROR test_error.py::test_error_during_setup[a] - assert 1 == 2",
+        "ERROR test_error.py::test_error_during_setup[b] - assert 1 == 2",
+        "ERROR test_error.py::test_error_during_setup[c] - assert 1 == 2",
+        ]
+    )
 
 def test_error_during_teardown(pytester):
     """
diff --git a/tests/test_yield.py b/tests/test_yield.py
new file mode 100644
index 0000000..b5daab9
--- /dev/null
+++ b/tests/test_yield.py
@@ -0,0 +1,14 @@
+
+def test_error_during_setup(pytester):
+    pytester.copy_example("examples/test_yield.py")
+    result = pytester.runpytest('-v', '-s')
+    result.assert_outcomes(passed=3)
+    result.stdout.re_match_lines(
+        [
+            ".*test_yield.a.*",
+            "setup",
+            ".*test_yield.b.*",
+            ".*test_yield.c.*",
+            "teardown",
+        ]
+    )
\ No newline at end of file

From c62f1887f2839ef86b93bc8d9e17cb4d62b1bb4b Mon Sep 17 00:00:00 2001
From: Brian Okken <1568356+okken@users.noreply.github.com>
Date: Mon, 16 Oct 2023 08:19:16 -0700
Subject: [PATCH 2/6] remove test_class.py

---
 examples/test_class.py | 14 --------------
 1 file changed, 14 deletions(-)
 delete mode 100644 examples/test_class.py

diff --git a/examples/test_class.py b/examples/test_class.py
deleted file mode 100644
index 8e7e075..0000000
--- a/examples/test_class.py
+++ /dev/null
@@ -1,14 +0,0 @@
-import pytest
-
-
-@pytest.fixture(scope='class')
-def a_fixture():
-    print('\nfixture setup')
-    yield 'some data from setup'
-    print('\nfixture teardown')
-
-
-class TestClass():
-    @pytest.mark.parametrize('x', ['a', 'b', 'c'])
-    def test_param(self, x, a_fixture):
-        assert a_fixture == 'some data from setup'
\ No newline at end of file

From 9a085a412d4b8edbd0083175772d83cadf2a54dd Mon Sep 17 00:00:00 2001
From: Brian Okken <1568356+okken@users.noreply.github.com>
Date: Wed, 18 Oct 2023 06:56:21 -0700
Subject: [PATCH 3/6] Update tests/test_yield.py

Co-authored-by: Blake Naccarato <blake.naccarato@gmail.com>
---
 tests/test_yield.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_yield.py b/tests/test_yield.py
index b5daab9..f64d854 100644
--- a/tests/test_yield.py
+++ b/tests/test_yield.py
@@ -1,7 +1,7 @@
 
 def test_error_during_setup(pytester):
     pytester.copy_example("examples/test_yield.py")
-    result = pytester.runpytest('-v', '-s')
+    result = pytester.runpytest('test_yield.py::test_yield', '-v', '-s')
     result.assert_outcomes(passed=3)
     result.stdout.re_match_lines(
         [

From d46b9763d1f986a82eb0d0a866ad7db1850a41ba Mon Sep 17 00:00:00 2001
From: Brian Okken <1568356+okken@users.noreply.github.com>
Date: Wed, 18 Oct 2023 11:51:49 -0700
Subject: [PATCH 4/6] mostly fix things up and add test. still need to update
 readme

---
 examples/test_error.py             |  1 -
 examples/test_marker_bad_params.py | 13 +++++++++
 examples/test_yield.py             | 36 +++++++++++++++++++++++
 src/pytest_param_scope/plugin.py   |  8 ++++--
 tests/test_error.py                | 17 +++++++++++
 tests/test_yield.py                | 46 ++++++++++++++++++++++++++++--
 6 files changed, 116 insertions(+), 5 deletions(-)
 create mode 100644 examples/test_marker_bad_params.py

diff --git a/examples/test_error.py b/examples/test_error.py
index 911e0a2..1615866 100644
--- a/examples/test_error.py
+++ b/examples/test_error.py
@@ -36,4 +36,3 @@ def test_error_during_teardown(x):
     """
     ...
 
-
diff --git a/examples/test_marker_bad_params.py b/examples/test_marker_bad_params.py
new file mode 100644
index 0000000..8dc6a1c
--- /dev/null
+++ b/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)`
+    """
+    ...
diff --git a/examples/test_yield.py b/examples/test_yield.py
index 50d9678..056af81 100644
--- a/examples/test_yield.py
+++ b/examples/test_yield.py
@@ -11,3 +11,39 @@ def setup_and_teardown():
 @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
+    """
+    ...
diff --git a/src/pytest_param_scope/plugin.py b/src/pytest_param_scope/plugin.py
index e0c4710..0e133f7 100644
--- a/src/pytest_param_scope/plugin.py
+++ b/src/pytest_param_scope/plugin.py
@@ -42,8 +42,10 @@ 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:
@@ -82,7 +84,9 @@ def param_scope(request):
             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()
 
diff --git a/tests/test_error.py b/tests/test_error.py
index ff4b3fc..57cfc27 100644
--- a/tests/test_error.py
+++ b/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.*"
+        ]
+    )
+
diff --git a/tests/test_yield.py b/tests/test_yield.py
index f64d854..bbd9df7 100644
--- a/tests/test_yield.py
+++ b/tests/test_yield.py
@@ -1,5 +1,5 @@
 
-def test_error_during_setup(pytester):
+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)
@@ -11,4 +11,46 @@ def test_error_during_setup(pytester):
             ".*test_yield.c.*",
             "teardown",
         ]
-    )
\ No newline at end of file
+    )
+
+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.*",
+        ]
+    )

From 3e70e4ef77f3cfa9bdfc06ba0ce64ec6678fc2d3 Mon Sep 17 00:00:00 2001
From: Brian Okken <1568356+okken@users.noreply.github.com>
Date: Wed, 18 Oct 2023 11:53:14 -0700
Subject: [PATCH 5/6] clean up coverage report

---
 .coveragerc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.coveragerc b/.coveragerc
index 75e08f8..a420cd5 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -4,4 +4,5 @@ branch = True
 [paths]
 source =
    src
-   .tox/*/site-packages
+   .tox/*/lib/*/site-packages
+

From 4b1d55412d7290723c81587fc485b67cfc6d4905 Mon Sep 17 00:00:00 2001
From: Brian Okken <1568356+okken@users.noreply.github.com>
Date: Wed, 18 Oct 2023 12:28:46 -0700
Subject: [PATCH 6/6] document yield

---
 README.md | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/README.md b/README.md
index 33d1d53..60325f5 100644
--- a/README.md
+++ b/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.