Feat: connector-builder-test-read CLI with custom components support (do not merge)

airbyte_cdk/cli/connector_builder_test_read/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""CLI for running test reads from the Airbyte Connector Builder.
+
+This CLI accepts a config file and an action to perform on the connector.
+
+Usage:
+    connector-builder-test-read [--config CONFIG] [--action ACTION]
+
+Options:
+    --action ACTION    The action to perform (e.g., test, validate).
+    --config CONFIG    Path to the config file.
+"""
+from airbyte_cdk.cli.connector_builder_test_read.run import run
+
+__all__ = [
+    "run",
+]

airbyte_cdk/connector_builder/connector_builder_handler.py → airbyte_cdk/cli/connector_builder_test_read/connector_builder_handler.py

@@ -16,6 +16,9 @@
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
+from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
+    INJECTED_MANIFEST,
+)
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
@@ -50,7 +53,7 @@ def get_limits(config: Mapping[str, Any]) -> TestReadLimits:
 
 
 def create_source(config: Mapping[str, Any], limits: TestReadLimits) -> ManifestDeclarativeSource:
-    manifest = config["__injected_declarative_manifest"]
+    manifest = config[INJECTED_MANIFEST]
     return ManifestDeclarativeSource(
         config=config,
         emit_connector_builder_messages=True,

airbyte_cdk/connector_builder/main.py → airbyte_cdk/cli/connector_builder_test_read/run.py

@@ -96,12 +96,17 @@ def handle_request(args: List[str]) -> str:
     ).decode()  # type: ignore[no-any-return] # Serializer.dump() always returns AirbyteMessage
 
 
-if __name__ == "__main__":
+def run(args: List[str] | None) -> None:
+    args = args or sys.argv[1:]
     try:
-        print(handle_request(sys.argv[1:]))
+        print(handle_request(args))
     except Exception as exc:
         error = AirbyteTracedException.from_exception(
             exc, message=f"Error handling request: {str(exc)}"
         )
         m = error.as_airbyte_message()
         print(orjson.dumps(AirbyteMessageSerializer.dump(m)).decode())
+
+
+if __name__ == "__main__":
+    run()

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -55,7 +55,11 @@
 
 
 class ManifestDeclarativeSource(DeclarativeSource):
-    """Declarative source defined by a manifest of low-code components that define source connector behavior"""
+    """Declarative source defined by a manifest of low-code components that define source connector behavior.
+
+    This class can also handle custom components if they are provided in the config using the
+    `__injected_components_py` config key.
+    """
 
     def __init__(
         self,

airbyte_cdk/test/declarative/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Declarative tests framework.
+
+This module provides fixtures and utilities for testing Airbyte sources and destinations
+in a declarative way.
+"""
+
+from airbyte_cdk.test.declarative.test_suites import (
+    ConnectorTestSuiteBase,
+    DestinationTestSuiteBase,
+    SourceTestSuiteBase,
+)
+
+__all__ = [
+    "ConnectorTestSuiteBase",
+    "DestinationTestSuiteBase",
+    "SourceTestSuiteBase",
+]

airbyte_cdk/test/declarative/connector_interfaces.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from typing import Protocol, Type
+
+
+class ConnectorInterface(Protocol):
+    """Protocol for Airbyte connectors."""
+
+    @classmethod
+    def launch(cls, args: list[str] | None): ...
+
+
+@dataclass
+class PythonWrapper:
+    """Wrapper for Python source and destination connectors."""
+
+    connector_class: Type["ConnectorInterface"]

airbyte_cdk/test/declarative/models/__init__.py

@@ -0,0 +1,7 @@
+from scenario import (
+    AcceptanceTestScenario,
+)
+
+__all__ = [
+    "AcceptanceTestScenario",
+]

airbyte_cdk/test/declarative/models/scenario.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Run acceptance tests in PyTest.
+
+These tests leverage the same `acceptance-test-config.yml` configuration files as the
+acceptance tests in CAT, but they run in PyTest instead of CAT. This allows us to run
+the acceptance tests in the same local environment as we are developing in, speeding
+up iteration cycles.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Literal
+
+import yaml
+from pydantic import BaseModel
+
+
+class AcceptanceTestScenario(BaseModel):
+    """Acceptance test instance, as a Pydantic model.
+
+    This class represents an acceptance test instance, which is a single test case
+    that can be run against a connector. It is used to deserialize and validate the
+    acceptance test configuration file.
+    """
+
+    class AcceptanceTestExpectRecords(BaseModel):
+        path: Path
+        exact_order: bool = False
+
+    class AcceptanceTestFileTypes(BaseModel):
+        skip_test: bool
+        bypass_reason: str
+
+    config_path: Path
+    configured_catalog_path: Path | None = None
+    timeout_seconds: int | None = None
+    expect_records: AcceptanceTestExpectRecords | None = None
+    file_types: AcceptanceTestFileTypes | None = None
+    status: Literal["succeed", "failed"] | None = None
+
+    @property
+    def expect_exception(self) -> bool:
+        return self.status and self.status == "failed"
+
+    @property
+    def instance_name(self) -> str:
+        return self.config_path.stem

airbyte_cdk/test/declarative/test_suites/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Declarative test suites.
+
+Here we have base classes for a robust set of declarative connector test suites.
+"""
+
+from airbyte_cdk.test.declarative.test_suites.connector_base import ConnectorTestSuiteBase
+from airbyte_cdk.test.declarative.test_suites.destination_base import DestinationTestSuiteBase
+from airbyte_cdk.test.declarative.test_suites.source_base import SourceTestSuiteBase
+
+__all__ = [
+    "ConnectorTestSuiteBase",
+    "DestinationTestSuiteBase",
+    "SourceTestSuiteBase",
+]

airbyte_cdk/test/declarative/test_suites/connector_base.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Base class for connector test suites."""
+
+from __future__ import annotations
+
+import abc
+from pathlib import Path
+from typing import Any, Literal
+
+import pytest
+import yaml
+from airbyte_connector_tester.job_runner import run_test_job
+from pydantic import BaseModel
+
+from airbyte_cdk import Connector
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    Type,
+)
+from airbyte_cdk.test import entrypoint_wrapper
+from airbyte_cdk.test.declarative.models import (
+    AcceptanceTestScenario,
+)
+
+ACCEPTANCE_TEST_CONFIG_PATH = Path("acceptance-test-config.yml")
+
+
+class ConnectorTestSuiteBase(abc.ABC):
+    """Base class for connector test suites."""
+
+    acceptance_test_file_path = Path("./acceptance-test-config.json")
+    """The path to the acceptance test config file.
+
+    By default, this is set to the `acceptance-test-config.json` file in
+    the root of the connector source directory.
+    """
+
+    connector_class: type[Connector]
+    """The connector class to test."""
+
+    # Public Methods - Subclasses may override these
+
+    @abc.abstractmethod
+    def new_connector(self, **kwargs: dict[str, Any]) -> Connector:
+        """Create a new connector instance.
+
+        By default, this returns a new instance of the connector class. Subclasses
+        may override this method to generate a dynamic connector instance.
+        """
+        return self.connector_factory()
+
+    # Internal Methods - We don't expect subclasses to override these
+
+    @classmethod
+    def _get_acceptance_tests(
+        category: str,
+        accept_test_config_path: Path = ACCEPTANCE_TEST_CONFIG_PATH,
+    ) -> list[AcceptanceTestScenario]:
+        all_tests_config = yaml.safe_load(accept_test_config_path.read_text())
+        if "acceptance_tests" not in all_tests_config:
+            raise ValueError(f"Acceptance tests config not found in {accept_test_config_path}")
+        if category not in all_tests_config["acceptance_tests"]:
+            return []
+        if "tests" not in all_tests_config["acceptance_tests"][category]:
+            raise ValueError(f"No tests found for category {category}")
+
+        return [
+            AcceptanceTestScenario.model_validate(test)
+            for test in all_tests_config["acceptance_tests"][category]["tests"]
+            if "iam_role" not in test["config_path"]
+        ]
+
+    # Test Definitions
+
+    @pytest.mark.parametrize(
+        "test_input,expected",
+        [
+            ("3+5", 8),
+            ("2+4", 6),
+            ("6*9", 54),
+        ],
+    )
+    def test_use_plugin_parametrized_test(
+        self,
+        test_input,
+        expected,
+    ):
+        assert eval(test_input) == expected
+
+    @pytest.mark.parametrize(
+        "instance",
+        self._get_acceptance_tests("connection"),
+        ids=lambda instance: instance.instance_name,
+    )
+    def test_check(
+        self,
+        instance: AcceptanceTestScenario,
+    ) -> None:
+        """Run `connection` acceptance tests."""
+        result: entrypoint_wrapper.EntrypointOutput = run_test_job(
+            self.new_connector(),
+            "check",
+            test_instance=instance,
+        )
+        conn_status_messages: list[AirbyteMessage] = [
+            msg for msg in result._messages if msg.type == Type.CONNECTION_STATUS
+        ]  # noqa: SLF001  # Non-public API
+        assert len(conn_status_messages) == 1, (
+            "Expected exactly one CONNECTION_STATUS message. Got: \n" + "\n".join(result._messages)
+        )

airbyte_cdk/test/declarative/test_suites/destination_base.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Base class for destination test suites."""
+
+from airbyte_connector_tester.connector_tests import ConnectorTestSuiteBase
+
+
+class DestinationTestSuiteBase(ConnectorTestSuiteBase):
+    """Base class for destination test suites.
+
+    This class provides a base set of functionality for testing destination connectors, and it
+    inherits all generic connector tests from the `ConnectorTestSuiteBase` class.
+    """