add support for 'spec' and tests without a scenario

Author: aaronsteers

airbyte_cdk/test/entrypoint_wrapper.py

@@ -82,6 +82,10 @@ def records(self) -> List[AirbyteMessage]:
     def state_messages(self) -> List[AirbyteMessage]:
         return self._get_message_by_types([Type.STATE])
 
+    @property
+    def spec_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.SPEC])
+
     @property
     def connection_status_messages(self) -> List[AirbyteMessage]:
         return self._get_message_by_types([Type.CONNECTION_STATUS])

airbyte_cdk/test/standard_tests/_job_runner.py

@@ -56,9 +56,9 @@ def spec(self, logger: logging.Logger) -> Any:
 
 def run_test_job(
     connector: IConnector | type[IConnector] | Callable[[], IConnector],
-    verb: Literal["read", "check", "discover"],
-    test_scenario: ConnectorTestScenario,
+    verb: Literal["spec", "read", "check", "discover"],
     *,
+    test_scenario: ConnectorTestScenario | None = None,
     catalog: ConfiguredAirbyteCatalog | dict[str, Any] | None = None,
 ) -> entrypoint_wrapper.EntrypointOutput:
     """Run a test scenario from provided CLI args and return the result."""
@@ -81,9 +81,9 @@ def run_test_job(
         )
 
     args: list[str] = [verb]
-    if test_scenario.config_path:
+    if test_scenario and test_scenario.config_path:
         args += ["--config", str(test_scenario.config_path)]
-    elif test_scenario.config_dict:
+    elif test_scenario and test_scenario.config_dict:
         config_path = (
             Path(tempfile.gettempdir()) / "airbyte-test" / f"temp_config_{uuid.uuid4().hex}.json"
         )
@@ -103,7 +103,7 @@ def run_test_job(
             )
             catalog_path.parent.mkdir(parents=True, exist_ok=True)
             catalog_path.write_text(orjson.dumps(catalog).decode())
-        elif test_scenario.configured_catalog_path:
+        elif test_scenario and test_scenario.configured_catalog_path:
             catalog_path = Path(test_scenario.configured_catalog_path)
 
         if catalog_path:
@@ -112,12 +112,18 @@ def run_test_job(
     # This is a bit of a hack because the source needs the catalog early.
     # Because it *also* can fail, we have to redundantly wrap it in a try/except block.
 
+    expect_exception = False
+    if test_scenario and test_scenario.expect_exception:
+        # If the test scenario expects an exception, we need to set the
+        # `expect_exception` flag to True.
+        expect_exception = True
+
     result: entrypoint_wrapper.EntrypointOutput = entrypoint_wrapper._run_command(  # noqa: SLF001  # Non-public API
         source=connector_obj,  # type: ignore [arg-type]
         args=args,
-        expecting_exception=test_scenario.expect_exception,
+        expecting_exception=False if not test_scenario else expect_exception,
     )
-    if result.errors and not test_scenario.expect_exception:
+    if result.errors and not expect_exception:
         raise AssertionError(
             f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
         )
@@ -132,7 +138,7 @@ def run_test_job(
             + "\n".join([str(msg) for msg in result.connection_status_messages])
             + _errors_to_str(result)
         )
-        if test_scenario.expect_exception:
+        if expect_exception:
             conn_status = result.connection_status_messages[0].connectionStatus
             assert conn_status, (
                 "Expected CONNECTION_STATUS message to be present. Got: \n"
@@ -146,7 +152,7 @@ def run_test_job(
         return result
 
     # For all other verbs, we assert check that an exception is raised (or not).
-    if test_scenario.expect_exception:
+    if expect_exception:
         if not result.errors:
             raise AssertionError("Expected exception but got none.")
 

airbyte_cdk/test/standard_tests/connector_base.py

@@ -89,7 +89,7 @@ def get_test_class_dir(cls) -> Path:
     @classmethod
     def create_connector(
         cls,
-        scenario: ConnectorTestScenario,
+        scenario: ConnectorTestScenario | None,
     ) -> IConnector:
         """Instantiate the connector class."""
         connector = cls.connector  # type: ignore
@@ -147,30 +147,33 @@ def get_scenarios(
         This has to be a separate function because pytest does not allow
         parametrization of fixtures with arguments from the test class itself.
         """
-        category = "connection"
+        categories = ["connection", "spec"]
         all_tests_config = yaml.safe_load(cls.acceptance_test_config_path.read_text())
         if "acceptance_tests" not in all_tests_config:
             raise ValueError(
                 f"Acceptance tests config not found in {cls.acceptance_test_config_path}."
                 f" Found only: {str(all_tests_config)}."
             )
 
-        if (
-            category not in all_tests_config["acceptance_tests"]
-            or "tests" not in all_tests_config["acceptance_tests"][category]
-        ):
-            return []
-
-        tests_scenarios = [
-            ConnectorTestScenario.model_validate(test)
-            for test in all_tests_config["acceptance_tests"][category]["tests"]
-            if "iam_role" not in test["config_path"]
-        ]
+        test_scenarios: list[ConnectorTestScenario] = []
+        for category in categories:
+            if (
+                category not in all_tests_config["acceptance_tests"]
+                or "tests" not in all_tests_config["acceptance_tests"][category]
+            ):
+                continue
+
+            test_scenarios.extend([
+                ConnectorTestScenario.model_validate(test)
+                for test in all_tests_config["acceptance_tests"][category]["tests"]
+                if "config_path" in test and "iam_role" not in test["config_path"]
+            ])
+
         connector_root = cls.get_connector_root_dir().absolute()
-        for test in tests_scenarios:
+        for test in test_scenarios:
             if test.config_path:
                 test.config_path = connector_root / test.config_path
             if test.configured_catalog_path:
                 test.configured_catalog_path = connector_root / test.configured_catalog_path
 
-        return tests_scenarios
+        return test_scenarios

airbyte_cdk/test/standard_tests/declarative_sources.py

@@ -64,7 +64,7 @@ def components_py_path(cls) -> Path | None:
     @classmethod
     def create_connector(
         cls,
-        scenario: ConnectorTestScenario,
+        scenario: ConnectorTestScenario | None,
     ) -> IConnector:
         """Create a connector scenario for the test suite.
 
@@ -73,9 +73,13 @@ def create_connector(
 
         Subclasses should not need to override this method.
         """
-        config: dict[str, Any] = scenario.get_config_dict()
-
         manifest_dict = yaml.safe_load(cls.manifest_yaml_path.read_text())
+        config = {
+            "__injected_manifest": manifest_dict,
+        }
+        if scenario:
+            config.update(scenario.get_config_dict())
+
         if cls.components_py_path and cls.components_py_path.exists():
             os.environ["AIRBYTE_ENABLE_UNSAFE_CODE"] = "true"
             config["__injected_components_py"] = cls.components_py_path.read_text()

airbyte_cdk/test/standard_tests/source_base.py

@@ -64,6 +64,26 @@ def test_discover(
             test_scenario=scenario,
         )
 
+    def test_spec(
+        self,
+        # scenario: ConnectorTestScenario,  # No inputs needed for spec test
+    ) -> None:
+        """Standard test for `spec`."""
+        result = run_test_job(
+            verb="spec",
+            test_scenario=None,
+            connector=self.create_connector(scenario=None),
+        )
+        if result.errors:
+            raise AssertionError(
+                f"Expected no errors but got {len(result.errors)}: \n"
+                + "\n".join([str(e) for e in result.errors])
+            )
+        assert len(result.spec_messages) == 1, (
+            "Expected exactly 1 spec message but got {len(result.spec_messages)}",
+            result.errors,
+        )
+
     def test_basic_read(
         self,
         scenario: ConnectorTestScenario,