airbyte_cdk.test.standard_tests.connector_base

  1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
  2"""Base class for connector test suites."""
  3
  4from __future__ import annotations
  5
  6import importlib
  7import os
  8from pathlib import Path
  9from typing import TYPE_CHECKING, cast
 10
 11from boltons.typeutils import classproperty
 12
 13from airbyte_cdk.test import entrypoint_wrapper
 14from airbyte_cdk.test.models import (
 15    ConnectorTestScenario,
 16)
 17from airbyte_cdk.test.standard_tests._job_runner import IConnector, run_test_job
 18from airbyte_cdk.test.standard_tests.docker_base import DockerConnectorTestSuite
 19
 20if TYPE_CHECKING:
 21    from collections.abc import Callable
 22
 23    from airbyte_cdk.test import entrypoint_wrapper
 24
 25
 26class ConnectorTestSuiteBase(DockerConnectorTestSuite):
 27    """Base class for Python connector test suites."""
 28
 29    connector: type[IConnector] | Callable[[], IConnector] | None  # type: ignore [reportRedeclaration]
 30    """The connector class or a factory function that returns an scenario of IConnector."""
 31
 32    @classproperty  # type: ignore [no-redef]
 33    def connector(cls) -> type[IConnector] | Callable[[], IConnector] | None:
 34        """Get the connector class for the test suite.
 35
 36        This assumes a python connector and should be overridden by subclasses to provide the
 37        specific connector class to be tested.
 38        """
 39        connector_root = cls.get_connector_root_dir()
 40        connector_name = cls.connector_name
 41
 42        expected_module_name = connector_name.replace("-", "_").lower()
 43        expected_class_name = connector_name.replace("-", "_").title().replace("_", "")
 44
 45        # dynamically import and get the connector class: <expected_module_name>.<expected_class_name>
 46
 47        cwd_snapshot = Path().absolute()
 48        os.chdir(connector_root)
 49
 50        # Dynamically import the module
 51        try:
 52            module = importlib.import_module(expected_module_name)
 53        except ModuleNotFoundError as e:
 54            raise ImportError(
 55                f"Could not import module '{expected_module_name}'. "
 56                "Please ensure you are running from within the connector's virtual environment, "
 57                "for instance by running `poetry run airbyte-cdk connector test` from the "
 58                "connector directory. If the issue persists, check that the connector "
 59                f"module matches the expected module name '{expected_module_name}' and that the "
 60                f"connector class matches the expected class name '{expected_class_name}'. "
 61                "Alternatively, you can run `airbyte-cdk image test` to run a subset of tests "
 62                "against the connector's image."
 63            ) from e
 64        finally:
 65            # Change back to the original working directory
 66            os.chdir(cwd_snapshot)
 67
 68        # Dynamically get the class from the module
 69        try:
 70            return cast(type[IConnector], getattr(module, expected_class_name))
 71        except AttributeError as e:
 72            # We did not find it based on our expectations, so let's check if we can find it
 73            # with a case-insensitive match.
 74            matching_class_name = next(
 75                (name for name in dir(module) if name.lower() == expected_class_name.lower()),
 76                None,
 77            )
 78            if not matching_class_name:
 79                raise ImportError(
 80                    f"Module '{expected_module_name}' does not have a class named '{expected_class_name}'."
 81                ) from e
 82            return cast(type[IConnector], getattr(module, matching_class_name))
 83
 84    @classmethod
 85    def create_connector(
 86        cls,
 87        scenario: ConnectorTestScenario | None,
 88    ) -> IConnector:
 89        """Instantiate the connector class."""
 90        connector = cls.connector  # type: ignore
 91        if connector:
 92            if callable(connector) or isinstance(connector, type):
 93                # If the connector is a class or factory function, instantiate it:
 94                return cast(IConnector, connector())  # type: ignore [redundant-cast]
 95
 96        # Otherwise, we can't instantiate the connector. Fail with a clear error message.
 97        raise NotImplementedError(
 98            "No connector class or connector factory function provided. "
 99            "Please provide a class or factory function in `cls.connector`, or "
100            "override `cls.create_connector()` to define a custom initialization process."
101        )
102
103    # Test Definitions
104
105    def test_check(
106        self,
107        scenario: ConnectorTestScenario,
108    ) -> None:
109        """Run `connection` acceptance tests."""
110        result: entrypoint_wrapper.EntrypointOutput = run_test_job(
111            self.create_connector(scenario),
112            "check",
113            test_scenario=scenario,
114            connector_root=self.get_connector_root_dir(),
115        )
116        assert len(result.connection_status_messages) == 1, (
117            f"Expected exactly one CONNECTION_STATUS message. "
118            "Got: {result.connection_status_messages!s}"
119        )