export testgen

pyproject.toml

@@ -62,6 +62,7 @@ test = [
     "types-paramiko>=2.7,<4",
     "types-python-dateutil>2,<3",
     "types-PyYAML>6,<7",
+    "pyinfra-testgen==0.1.1",
 ]
 
 docs = [

tests/test_facts.py

@@ -4,12 +4,14 @@
 from os import listdir, path
 from unittest import TestCase
 
+from testgen import TestGenerator
+
 from pyinfra.api import StringCommand
 from pyinfra.api.facts import ShortFactBase
 from pyinfra.context import ctx_host, ctx_state
 from pyinfra_cli.util import json_encode
 
-from .util import FakeState, YamlTest, create_host, get_command_string
+from .util import FakeState, create_host, get_command_string
 
 # show full diff on json
 TestCase.maxDiff = None
@@ -32,16 +34,19 @@ def make_fact_tests(folder_name):
     module = import_module("pyinfra.facts.{0}".format(module_name))
     fact = getattr(module, fact_name)()
 
-    class TestTests(TestCase, metaclass=YamlTest):
-        yaml_test_dir = path.join("tests", "facts", folder_name)
-        yaml_test_prefix = "test_{0}_".format(fact.name)
-
+    class TestTests(
+        TestCase,
+        metaclass=TestGenerator,
+        tests_dir=path.join("tests", "facts", folder_name),
+        test_prefix="test_{0}_".format(fact.name),
+        test_method="_test",
+    ):
         @classmethod
         def setUpClass(cls):
             # Create a global fake state that attach to context state
             cls.state = FakeState()
 
-        def yaml_test_function(self, test_name, test_data, fact=fact):
+        def _test(self, test_name, test_data, fact=fact):
             host = create_host(self.state, facts=test_data.get("facts", {}))
             with ctx_state.use(self.state):
                 with ctx_host.use(host):

tests/test_operations.py

@@ -6,11 +6,13 @@
 from unittest import TestCase
 from unittest.mock import patch
 
+from testgen import TestGenerator
+
 from pyinfra.api import FileDownloadCommand, FileUploadCommand, FunctionCommand, StringCommand
 from pyinfra.context import ctx_host, ctx_state
 from pyinfra_cli.util import json_encode
 
-from .util import FakeState, YamlTest, create_host, get_command_string, parse_value, patch_files
+from .util import FakeState, create_host, get_command_string, parse_value, patch_files
 
 PLATFORM_NAME = platform.system()
 
@@ -92,16 +94,19 @@ def make_operation_tests(arg):
     # Generate a test class
     @patch("pyinfra.operations.files.get_timestamp", lambda: "a-timestamp")
     @patch("pyinfra.operations.util.files.get_timestamp", lambda: "a-timestamp")
-    class TestTests(TestCase, metaclass=YamlTest):
-        yaml_test_dir = path.join("tests", "operations", arg)
-        yaml_test_prefix = "test_{0}_{1}_".format(module_name, op_name)
-
+    class TestTests(
+        TestCase,
+        metaclass=TestGenerator,
+        tests_dir=path.join("tests", "operations", arg),
+        test_prefix="test_{0}_{1}_".format(module_name, op_name),
+        test_method="_test",
+    ):
         @classmethod
         def setUpClass(cls):
             # Create a global fake state that attach to context state
             cls.state = FakeState()
 
-        def yaml_test_function(self, test_name, test_data):
+        def _test(self, test_name, test_data):
             if (
                 "require_platform" in test_data
                 and PLATFORM_NAME not in test_data["require_platform"]

tests/util.py

@@ -8,8 +8,6 @@
 from pathlib import Path
 from unittest.mock import patch
 
-import yaml
-
 from pyinfra.api import Config, Inventory
 from pyinfra.api.util import get_kwargs_str
 
@@ -479,30 +477,3 @@ def create_host(state, name=None, facts=None, data=None):
         real_facts[name] = fact_data
 
     return FakeHost(state, name, facts=real_facts, data=data)
-
-
-class YamlTest(type):
-    def __new__(cls, name, bases, attrs):
-        test_suffixes = {".yaml", ".yml", ".json"}
-
-        tests_dir = Path(attrs["yaml_test_dir"])
-
-        test_files = [f for f in tests_dir.iterdir() if f.suffix in test_suffixes]
-
-        test_prefix = attrs.get("yaml_test_prefix", "test_")
-
-        def gen_test(test_name, test_file):
-            def test(self):
-                test_data = yaml.safe_load(test_file.open(encoding="utf-8").read())
-                self.yaml_test_function(test_name, test_data)
-
-            return test
-
-        # Loop them and create class methods to call the yaml_test_function
-        for test_file in test_files:
-            test_name = test_file.stem
-            # Attach the method
-            method_name = "{0}{1}".format(test_prefix, test_name)
-            attrs[method_name] = gen_test(test_name, test_file)
-
-        return type.__new__(cls, name, bases, attrs)