From 217429a12612aeaffb7336593aa9a4eb61cd2017 Mon Sep 17 00:00:00 2001
From: Kyle Altendorf <sda@fstab.net>
Date: Fri, 12 May 2023 13:58:45 -0400
Subject: [PATCH] add `datacases()` and `named_datacases()` (#15265)

* add `datacases()` and `named_datacases()`

* correct DataCasesProtocol

* add back the tests for testing the test utilities
---
 pytest.ini                    |  4 ++++
 tests/util/misc.py            | 33 ++++++++++++++++++++++++++++--
 tests/util/test_tests_misc.py | 38 +++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 tests/util/test_tests_misc.py

diff --git a/pytest.ini b/pytest.ini
index a02577aaa5..87e2b15214 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -9,6 +9,10 @@ asyncio_mode = strict
 markers =
     benchmark
     data_layer: Mark as a data layer related test.
+    test_mark_a1: used in testing test utilities
+    test_mark_a2: used in testing test utilities
+    test_mark_b1: used in testing test utilities
+    test_mark_b2: used in testing test utilities
 testpaths = tests
 filterwarnings =
     error
diff --git a/tests/util/misc.py b/tests/util/misc.py
index 344456b593..b6511a0964 100644
--- a/tests/util/misc.py
+++ b/tests/util/misc.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import contextlib
 import dataclasses
 import enum
+import functools
 import gc
 import math
 import os
@@ -13,10 +14,10 @@ from statistics import mean
 from textwrap import dedent
 from time import thread_time
 from types import TracebackType
-from typing import Any, Callable, Iterator, List, Optional, Type, Union
+from typing import Any, Callable, Collection, Iterator, List, Optional, Type, Union
 
 import pytest
-from typing_extensions import final
+from typing_extensions import Protocol, final
 
 from tests.core.data_layer.util import ChiaRoot
 
@@ -303,3 +304,31 @@ def closing_chia_root_popen(chia_root: ChiaRoot, args: List[str]) -> Iterator[su
                 process.wait(timeout=10)
             except subprocess.TimeoutExpired:
                 process.kill()
+
+
+# https://github.com/pytest-dev/pytest/blob/7.3.1/src/_pytest/mark/__init__.py#L45
+Marks = Union[pytest.MarkDecorator, Collection[Union[pytest.MarkDecorator, pytest.Mark]]]
+
+
+class DataCase(Protocol):
+    marks: Marks
+
+    @property
+    def id(self) -> str:
+        ...
+
+
+def datacases(*cases: DataCase, _name: str = "case") -> pytest.MarkDecorator:
+    return pytest.mark.parametrize(
+        argnames=_name,
+        argvalues=[pytest.param(case, id=case.id, marks=case.marks) for case in cases],
+    )
+
+
+class DataCasesDecorator(Protocol):
+    def __call__(self, *cases: DataCase, _name: str = "case") -> pytest.MarkDecorator:
+        ...
+
+
+def named_datacases(name: str) -> DataCasesDecorator:
+    return functools.partial(datacases, _name=name)
diff --git a/tests/util/test_tests_misc.py b/tests/util/test_tests_misc.py
new file mode 100644
index 0000000000..01836b3549
--- /dev/null
+++ b/tests/util/test_tests_misc.py
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pytest
+
+from tests.util.misc import Marks, datacases, named_datacases
+
+
+@dataclass
+class DataCase:
+    id: str
+    marks: Marks
+
+
+sample_cases = [
+    DataCase(id="id_a", marks=[pytest.mark.test_mark_a1, pytest.mark.test_mark_a2]),
+    DataCase(id="id_b", marks=[pytest.mark.test_mark_b1, pytest.mark.test_mark_b2]),
+]
+
+
+def sample_result(name: str) -> pytest.MarkDecorator:
+    return pytest.mark.parametrize(
+        argnames=name,
+        argvalues=[pytest.param(case, id=case.id, marks=case.marks) for case in sample_cases],
+    )
+
+
+def test_datacases() -> None:
+    result = datacases(*sample_cases)
+
+    assert result == sample_result(name="case")
+
+
+def test_named_datacases() -> None:
+    result = named_datacases("Sharrilanda")(*sample_cases)
+
+    assert result == sample_result(name="Sharrilanda")