fix lint issues

Author: aaronsteers

airbyte_cdk/test/declarative/test_suites/connector_base.py

@@ -4,32 +4,23 @@
 from __future__ import annotations
 
 import abc
-import functools
-import inspect
-import sys
 from pathlib import Path
-from typing import Any, Callable, Literal
+from typing import Any, Literal
 
 import pytest
 import yaml
-from pydantic import BaseModel
-from typing_extensions import override
+from boltons.typeutils import classproperty
 
 from airbyte_cdk import Connector
 from airbyte_cdk.models import (
     AirbyteMessage,
     Type,
 )
-from airbyte_cdk.sources import Source
-from airbyte_cdk.sources.declarative.declarative_source import (
-    AbstractSource,
-    DeclarativeSource,
-)
 from airbyte_cdk.test import entrypoint_wrapper
 from airbyte_cdk.test.declarative.models import (
     ConnectorTestScenario,
 )
-from airbyte_cdk.test.declarative.utils.job_runner import run_test_job
+from airbyte_cdk.test.declarative.utils.job_runner import IConnector, run_test_job
 
 ACCEPTANCE_TEST_CONFIG = "acceptance-test-config.yml"
 
@@ -49,7 +40,7 @@ class RunnableConnector(abc.ABC):
     def launch(cls, args: list[str] | None) -> None: ...
 
 
-def generate_tests(metafunc) -> None:
+def generate_tests(metafunc: pytest.Metafunc) -> None:
     """
     A helper for pytest_generate_tests hook.
 
@@ -111,8 +102,9 @@ class ConnectorTestSuiteBase(abc.ABC):
 
     @classmethod
     def create_connector(
-        cls, scenario: ConnectorTestScenario
-    ) -> Source | AbstractSource | DeclarativeSource | RunnableConnector:
+        cls,
+        scenario: ConnectorTestScenario,
+    ) -> IConnector:
         """Instantiate the connector class."""
         raise NotImplementedError("Subclasses must implement this method.")
 
@@ -121,7 +113,7 @@ def run_test_scenario(
         verb: Literal["read", "check", "discover"],
         test_scenario: ConnectorTestScenario,
         *,
-        catalog: dict | None = None,
+        catalog: dict[str, Any] | None = None,
     ) -> entrypoint_wrapper.EntrypointOutput:
         """Run a test job from provided CLI args and return the result."""
         return run_test_job(
@@ -146,11 +138,10 @@ def test_check(
             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)
+            f"Expected exactly one CONNECTION_STATUS message. Got: {result._messages}"
         )
 
-    @classmethod
-    @property
+    @classproperty
     def acceptance_test_config_path(self) -> Path:
         """Get the path to the acceptance test config file.
 

airbyte_cdk/test/declarative/test_suites/source_base.py

@@ -45,8 +45,10 @@ def test_check(
         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)
+        num_status_messages = len(conn_status_messages)
+        assert num_status_messages == 1, (
+            f"Expected exactly one CONNECTION_STATUS message. Got {num_status_messages}: \n"
+            + "\n".join([str(m) for m in result._messages])
         )
 
     def test_basic_read(
@@ -70,7 +72,7 @@ def test_basic_read(
                     sync_mode=SyncMode.full_refresh,
                     destination_sync_mode=DestinationSyncMode.append_dedup,
                 )
-                for stream in discover_result.catalog.catalog.streams
+                for stream in discover_result.catalog.catalog.streams # type: ignore [reportOptionalMemberAccess]
             ]
         )
         result = run_test_job(
@@ -102,8 +104,8 @@ def test_fail_with_bad_catalog(
                         },
                         supported_sync_modes=[SyncMode.full_refresh],
                     ),
-                    sync_mode="INVALID",
-                    destination_sync_mode="INVALID",
+                    sync_mode="INVALID", # type: ignore [reportArgumentType]
+                    destination_sync_mode="INVALID", # type: ignore [reportArgumentType]
                 )
             ]
         )

airbyte_cdk/test/declarative/utils/job_runner.py

@@ -1,53 +1,56 @@
 import tempfile
 import uuid
+from dataclasses import asdict
 from pathlib import Path
 from typing import Any, Callable, Literal
 
 import orjson
+from typing_extensions import Protocol, runtime_checkable
 
-from airbyte_cdk import Connector
 from airbyte_cdk.models import (
+    ConfiguredAirbyteCatalog,
     Status,
 )
-from airbyte_cdk.sources.abstract_source import AbstractSource
-from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
 from airbyte_cdk.test import entrypoint_wrapper
 from airbyte_cdk.test.declarative.models import (
     ConnectorTestScenario,
 )
 
 
+@runtime_checkable
+class IConnector(Protocol):
+    """A connector that can be run in a test scenario."""
+
+    def launch(self, args: list[str] | None) -> None:
+        """Launch the connector with the given arguments."""
+        ...
+
+
 def run_test_job(
-    connector: Connector | type[Connector] | Callable[[], Connector],
+    connector: IConnector | type[IConnector] | Callable[[], IConnector],
     verb: Literal["read", "check", "discover"],
     test_instance: ConnectorTestScenario,
     *,
-    catalog: dict[str, Any] | None = None,
+    catalog: ConfiguredAirbyteCatalog | dict[str, Any] | None = None,
 ) -> entrypoint_wrapper.EntrypointOutput:
     """Run a test job from provided CLI args and return the result."""
     if not connector:
         raise ValueError("Connector is required")
 
-    connector_obj: Connector
-    if isinstance(connector, type):
+    if catalog and isinstance(catalog, ConfiguredAirbyteCatalog):
+        # Convert the catalog to a dict if it's already a ConfiguredAirbyteCatalog.
+        catalog = asdict(catalog)
+
+    connector_obj: IConnector
+    if isinstance(connector, type) or callable(connector):
+        # If the connector is a class or a factory lambda, instantiate it.
         connector_obj = connector()
-    elif isinstance(connector, Connector):
-        connector_obj = connector
-    elif isinstance(connector, DeclarativeSource | AbstractSource):
+    elif isinstance(connector, IConnector):
         connector_obj = connector
-    elif isinstance(connector, Callable):
-        try:
-            connector_obj = connector()
-        except Exception as ex:
-            if not test_instance.expect_exception:
-                raise
-
-            return entrypoint_wrapper.EntrypointOutput(
-                messages=[],
-                uncaught_exception=ex,
-            )
     else:
-        raise ValueError(f"Invalid source type: {type(connector)}")
+        raise ValueError(
+            f"Invalid connector input: {type(connector)}",
+        )
 
     args: list[str] = [verb]
     if test_instance.config_path:
@@ -82,34 +85,40 @@ def run_test_job(
     # Because it *also* can fail, we have ot redundantly wrap it in a try/except block.
 
     result: entrypoint_wrapper.EntrypointOutput = entrypoint_wrapper._run_command(  # noqa: SLF001  # Non-public API
-        source=connector_obj,
+        source=connector_obj,  # type: ignore [reportArgumentType]
         args=args,
         expecting_exception=test_instance.expect_exception,
     )
     if result.errors and not test_instance.expect_exception:
         raise AssertionError(
             "\n\n".join(
-                [str(err.trace.error).replace("\\n", "\n") for err in result.errors],
+                [str(err.trace.error).replace("\\n", "\n") for err in result.errors if err.trace],
             )
         )
 
     if verb == "check":
         # Check is expected to fail gracefully without an exception.
         # Instead, we assert that we have a CONNECTION_STATUS message with
         # a failure status.
-        assert not result.errors, "Expected no errors from check. Got:\n" + "\n".join(
-            [str(error) for error in result.errors]
-        )
+        assert not result.errors, "Expected no errors from check. Got:\n" + "\n".join([
+            str(error) for error in result.errors
+        ])
         assert len(result.connection_status_messages) == 1, (
             "Expected exactly one CONNECTION_STATUS message. Got "
             f"{len(result.connection_status_messages)}:\n"
             + "\n".join([str(msg) for msg in result.connection_status_messages])
         )
         if test_instance.expect_exception:
-            assert result.connection_status_messages[0].connectionStatus.status == Status.FAILED, (
+            conn_status = result.connection_status_messages[0].connectionStatus
+            assert conn_status, (
+                "Expected CONNECTION_STATUS message to be present. Got: \n"
+                + "\n".join([str(msg) for msg in result.connection_status_messages])
+            )
+            assert conn_status.status == Status.FAILED, (
                 "Expected CONNECTION_STATUS message to be FAILED. Got: \n"
                 + "\n".join([str(msg) for msg in result.connection_status_messages])
             )
+
         return result
 
     # For all other verbs, we assert check that an exception is raised (or not).
@@ -121,7 +130,14 @@ def run_test_job(
     if result.errors:
         raise AssertionError(
             "\n\n".join(
-                [str(err.trace.error).replace("\\n", "\n") for err in result.errors],
+                [
+                    str(err.trace.error).replace(
+                        "\\n",
+                        "\n",
+                    )
+                    for err in result.errors
+                    if err.trace
+                ],
             )
         )
 

airbyte_cdk/test/pytest_config/plugin.py

@@ -1,12 +1,12 @@
 """Global pytest configuration for the Airbyte CDK tests."""
 
 from pathlib import Path
-from typing import Optional
+from typing import cast
 
 import pytest
 
 
-def pytest_collect_file(parent: Optional[pytest.Module], path: Path) -> pytest.Module | None:
+def pytest_collect_file(parent: pytest.Module | None, path: Path) -> pytest.Module | None:
     """Collect test files based on their names."""
     if path.name == "test_connector.py":
         return cast(pytest.Module, pytest.Module.from_parent(parent, path=path))
@@ -41,6 +41,6 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
     print(f"Setting up test: {item.name}")
 
 
-def pytest_runtest_teardown(item: pytest.Item, nextitem: Optional[pytest.Item]) -> None:
+def pytest_runtest_teardown(item: pytest.Item, nextitem: pytest.Item | None) -> None:
     # This hook is called after each test function is executed
     print(f"Tearing down test: {item.name}")

poetry.lock

@@ -83,6 +83,7 @@ sqlalchemy = {version = "^2.0,!=2.0.36", optional = true }
 xmltodict = ">=0.13,<0.15"
 anyascii = "^0.3.2"
 whenever = "^0.6.16"
+boltons = "^25.0.0"
 
 [tool.poetry.group.dev.dependencies]
 freezegun = "*"