refactor: Break up TapTestClassFactory._annotate_test_class into simp…

…ler methods
meltano · Nov 22, 2023 · 6884436 · 6884436
1 parent 4cc52db
commit 6884436
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 28 deletions.
diff --git a/singer_sdk/testing/factory.py b/singer_sdk/testing/factory.py
@@ -8,14 +8,16 @@
 from .config import SuiteConfig
 from .runners import TapTestRunner, TargetTestRunner
 from .suites import (
+    TestSuite,
     tap_stream_attribute_tests,
     tap_stream_tests,
     tap_tests,
     target_tests,
 )
 
 if t.TYPE_CHECKING:
-    from singer_sdk import Tap, Target
+    from singer_sdk import Stream, Tap, Target
+    from singer_sdk.testing.templates import StreamTestTemplate, TapTestTemplate
 
 
 class BaseTestClass:
@@ -145,7 +147,7 @@ def runner(self) -> TapTestRunner | TargetTestRunner:
         return TapTestClass
 
     # TODO: Refactor this. It's too long and nested.
-    def _annotate_test_class(  # noqa: C901
+    def _annotate_test_class(
         self,
         empty_test_class: type[BaseTestClass],
         test_suites: list,
@@ -163,33 +165,13 @@ def _annotate_test_class(  # noqa: C901
         """
         for suite in test_suites:
             if suite.kind == "tap":
-                for test_class in suite.tests:
-                    test = test_class()
-                    test_name = f"test_{suite.kind}_{test.name}"
-                    setattr(empty_test_class, test_name, test.run)
+                self._with_tap_tests(empty_test_class, suite)
 
             if suite.kind in {"tap_stream", "tap_stream_attribute"}:
                 streams = list(test_runner.new_tap().streams.values())
 
                 if suite.kind == "tap_stream":
-                    params = [
-                        {
-                            "stream": stream,
-                        }
-                        for stream in streams
-                    ]
-                    param_ids = [stream.name for stream in streams]
-
-                    for test_class in suite.tests:
-                        test = test_class()
-                        test_name = f"test_{suite.kind}_{test.name}"
-                        setattr(
-                            empty_test_class,
-                            test_name,
-                            test.run,
-                        )
-                        empty_test_class.params[test_name] = params
-                        empty_test_class.param_ids[test_name] = param_ids
+                    self._with_stream_tests(empty_test_class, suite, streams)
 
                 if suite.kind == "tap_stream_attribute":
                     for test_class in suite.tests:
@@ -238,6 +220,41 @@ def _annotate_test_class(  # noqa: C901
 
         return empty_test_class
 
+    def _with_tap_tests(
+        self,
+        empty_test_class: type[BaseTestClass],
+        suite: TestSuite[TapTestTemplate],
+    ) -> None:
+        for test_class in suite.tests:
+            test = test_class()
+            test_name = f"test_{suite.kind}_{test.name}"
+            setattr(empty_test_class, test_name, test.run)
+
+    def _with_stream_tests(
+        self,
+        empty_test_class: type[BaseTestClass],
+        suite: TestSuite[StreamTestTemplate],
+        streams: list[Stream],
+    ) -> None:
+        params = [
+            {
+                "stream": stream,
+            }
+            for stream in streams
+        ]
+        param_ids = [stream.name for stream in streams]
+
+        for test_class in suite.tests:
+            test = test_class()
+            test_name = f"test_{suite.kind}_{test.name}"
+            setattr(
+                empty_test_class,
+                test_name,
+                test.run,
+            )
+            empty_test_class.params[test_name] = params
+            empty_test_class.param_ids[test_name] = param_ids
+
 
 class TargetTestClassFactory:
     """Factory for Target Test Classes."""

diff --git a/singer_sdk/testing/suites.py b/singer_sdk/testing/suites.py
@@ -42,17 +42,17 @@
     TargetSchemaUpdates,
     TargetSpecialCharsInAttributes,
 )
+from .templates import TestTemplate
 
-if t.TYPE_CHECKING:
-    from .templates import TapTestTemplate, TargetTestTemplate, TestTemplate
+T = t.TypeVar("T", bound=TestTemplate)
 
 
 @dataclass
-class TestSuite:
+class TestSuite(t.Generic[T]):
     """Test Suite container class."""
 
     kind: str
-    tests: list[type[TestTemplate] | type[TapTestTemplate] | type[TargetTestTemplate]]
+    tests: list[type[T]]
 
 
 # Tap Test Suites