airbyte-cdk 6.45.0.post24.dev14387467928__tar.gz → 6.45.1.post44.dev14417290834__tar.gz

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 6.45.0.post24.dev14387467928
+Version: 6.45.1.post44.dev14417290834
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://airbyte.com
 License: MIT
@@ -26,6 +26,7 @@ Requires-Dist: airbyte-protocol-models-dataclasses (>=0.14,<0.15)
 Requires-Dist: anyascii (>=0.3.2,<0.4.0)
 Requires-Dist: avro (>=1.11.2,<1.13.0) ; extra == "file-based"
 Requires-Dist: backoff
+Requires-Dist: boltons (>=25.0.0,<26.0.0)
 Requires-Dist: cachetools
 Requires-Dist: cohere (==4.21) ; extra == "vector-db-based"
 Requires-Dist: cryptography (>=44.0.0,<45.0.0)

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -35,8 +35,10 @@ MAX_RECORDS_KEY = "max_records"
 MAX_STREAMS_KEY = "max_streams"
 
 
-@dataclass
+@dataclass(kw_only=True)
 class TestLimits:
+    __test__: bool = False  # Prevent pytest from treating this as a test case, despite its name
+
     max_records: int = field(default=DEFAULT_MAXIMUM_RECORDS)
     max_pages_per_slice: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
     max_slices: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
@@ -51,7 +53,12 @@ def get_limits(config: Mapping[str, Any]) -> TestLimits:
     max_slices = command_config.get(MAX_SLICES_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_SLICES
     max_records = command_config.get(MAX_RECORDS_KEY) or DEFAULT_MAXIMUM_RECORDS
     max_streams = command_config.get(MAX_STREAMS_KEY) or DEFAULT_MAXIMUM_STREAMS
-    return TestLimits(max_records, max_pages_per_slice, max_slices, max_streams)
+    return TestLimits(
+        max_records=max_records,
+        max_pages_per_slice=max_pages_per_slice,
+        max_slices=max_slices,
+        max_streams=max_streams,
+    )
 
 
 def create_source(config: Mapping[str, Any], limits: TestLimits) -> ManifestDeclarativeSource:
@@ -79,7 +86,9 @@ def read_stream(
 ) -> AirbyteMessage:
     try:
         test_read_handler = TestReader(
-            limits.max_pages_per_slice, limits.max_slices, limits.max_records
+            max_pages_per_slice=limits.max_pages_per_slice,
+            max_slices=limits.max_slices,
+            max_record_limit=limits.max_records,
         )
         # The connector builder only supports a single stream
         stream_name = configured_catalog.streams[0].stream.name

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/connector_builder/main.py

@@ -78,9 +78,9 @@ def handle_connector_builder_request(
     if command == "resolve_manifest":
         return resolve_manifest(source)
     elif command == "test_read":
-        assert (
-            catalog is not None
-        ), "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
+        assert catalog is not None, (
+            "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
+        )
         return read_stream(source, config, catalog, state, limits)
     elif command == "full_resolve_manifest":
         return full_resolve_manifest(source, limits)

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/connector_builder/test_reader/reader.py

@@ -66,6 +66,8 @@ class TestReader:
 
     """
 
+    __test__: bool = False  # Prevent pytest from treating this as a test case, despite its name
+
     logger = logging.getLogger("airbyte.connector-builder")
 
     def __init__(

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/sources/concurrent_source/concurrent_source.py

@@ -49,9 +49,9 @@ class ConcurrentSource:
         too_many_generator = (
             not is_single_threaded and initial_number_of_partitions_to_generate >= num_workers
         )
-        assert (
-            not too_many_generator
-        ), "It is required to have more workers than threads generating partitions"
+        assert not too_many_generator, (
+            "It is required to have more workers than threads generating partitions"
+        )
         threadpool = ThreadPoolManager(
             concurrent.futures.ThreadPoolExecutor(
                 max_workers=num_workers, thread_name_prefix="workerpool"

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/sources/declarative/interpolation/filters.py

@@ -4,9 +4,10 @@
 
 import base64
 import hashlib
+import hmac as hmac_lib
 import json
 import re
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 
 def hash(value: Any, hash_type: str = "md5", salt: Optional[str] = None) -> str:
@@ -135,5 +136,51 @@ def regex_search(value: str, regex: str) -> str:
     return ""
 
 
-_filters_list = [hash, base64encode, base64decode, base64binascii_decode, string, regex_search]
+def hmac(value: Any, key: str, hash_type: str = "sha256") -> str:
+    """
+    Implementation of a custom Jinja2 hmac filter with SHA-256 support.
+
+    This filter creates a Hash-based Message Authentication Code (HMAC) using a cryptographic
+    hash function and a secret key. Currently only supports SHA-256, and returns hexdigest of the signature.
+
+    Example usage in a low code connector:
+
+    auth_headers:
+      $ref: "#/definitions/base_auth"
+      $parameters:
+        signature: "{{ 'message_to_sign' | hmac('my_secret_key') }}"
+
+    :param value: The message to be authenticated
+    :param key: The secret key for the HMAC
+    :param hash_type: Hash algorithm to use (default: sha256)
+    :return: HMAC digest as a hexadecimal string
+    """
+    # Define allowed hash functions
+    ALLOWED_HASH_TYPES: Dict[str, Any] = {
+        "sha256": hashlib.sha256,
+    }
+
+    if hash_type not in ALLOWED_HASH_TYPES:
+        raise ValueError(
+            f"Hash type '{hash_type}' is not allowed. Allowed types: {', '.join(ALLOWED_HASH_TYPES.keys())}"
+        )
+
+    hmac_obj = hmac_lib.new(
+        key=str(key).encode("utf-8"),
+        msg=str(value).encode("utf-8"),
+        digestmod=ALLOWED_HASH_TYPES[hash_type],
+    )
+
+    return hmac_obj.hexdigest()
+
+
+_filters_list = [
+    hash,
+    base64encode,
+    base64decode,
+    base64binascii_decode,
+    string,
+    regex_search,
+    hmac,
+]
 filters = {f.__name__: f for f in _filters_list}

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/sources/file_based/file_based_source.py

@@ -282,9 +282,9 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
                     and hasattr(self, "_concurrency_level")
                     and self._concurrency_level is not None
                 ):
-                    assert (
-                        state_manager is not None
-                    ), "No ConnectorStateManager was created, but it is required for incremental syncs. This is unexpected. Please contact Support."
+                    assert state_manager is not None, (
+                        "No ConnectorStateManager was created, but it is required for incremental syncs. This is unexpected. Please contact Support."
+                    )
 
                     cursor = self.cursor_cls(
                         stream_config,

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/sources/file_based/file_types/avro_parser.py

@@ -154,7 +154,7 @@ class AvroParser(FileTypeParser):
                 # For example: ^-?\d{1,5}(?:\.\d{1,3})?$ would accept 12345.123 and 123456.12345 would  be rejected
                 return {
                     "type": "string",
-                    "pattern": f"^-?\\d{{{1,max_whole_number_range}}}(?:\\.\\d{1,decimal_range})?$",
+                    "pattern": f"^-?\\d{{{1, max_whole_number_range}}}(?:\\.\\d{1, decimal_range})?$",
                 }
             elif "logicalType" in avro_field:
                 if avro_field["logicalType"] not in AVRO_LOGICAL_TYPE_TO_JSON:

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/sources/file_based/stream/concurrent/adapters.py

@@ -284,9 +284,9 @@ class FileBasedStreamPartition(Partition):
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         if self._slice is None:
             return None
-        assert (
-            len(self._slice["files"]) == 1
-        ), f"Expected 1 file per partition but got {len(self._slice['files'])} for stream {self.stream_name()}"
+        assert len(self._slice["files"]) == 1, (
+            f"Expected 1 file per partition but got {len(self._slice['files'])} for stream {self.stream_name()}"
+        )
         file = self._slice["files"][0]
         return {"files": [file]}
 

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/sql/shared/sql_processor.py

@@ -326,9 +326,9 @@ class SqlProcessorBase(abc.ABC):
 
         if DEBUG_MODE:
             found_schemas = schemas_list
-            assert (
-                schema_name in found_schemas
-            ), f"Schema {schema_name} was not created. Found: {found_schemas}"
+            assert schema_name in found_schemas, (
+                f"Schema {schema_name} was not created. Found: {found_schemas}"
+            )
 
     def _quote_identifier(self, identifier: str) -> str:
         """Return the given identifier, quoted."""
@@ -617,10 +617,10 @@ class SqlProcessorBase(abc.ABC):
         self._execute_sql(
             f"""
             INSERT INTO {self._fully_qualified(final_table_name)} (
-            {f',{nl}  '.join(columns)}
+            {f",{nl}  ".join(columns)}
             )
             SELECT
-            {f',{nl}  '.join(columns)}
+            {f",{nl}  ".join(columns)}
             FROM {self._fully_qualified(temp_table_name)}
             """,
         )
@@ -645,8 +645,7 @@ class SqlProcessorBase(abc.ABC):
         deletion_name = f"{final_table_name}_deleteme"
         commands = "\n".join(
             [
-                f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME "
-                f"TO {deletion_name};",
+                f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME TO {deletion_name};",
                 f"ALTER TABLE {self._fully_qualified(temp_table_name)} RENAME "
                 f"TO {final_table_name};",
                 f"DROP TABLE {self._fully_qualified(deletion_name)};",
@@ -686,10 +685,10 @@ class SqlProcessorBase(abc.ABC):
                 {set_clause}
             WHEN NOT MATCHED THEN INSERT
             (
-                {f',{nl}    '.join(columns)}
+                {f",{nl}    ".join(columns)}
             )
             VALUES (
-                tmp.{f',{nl}    tmp.'.join(columns)}
+                tmp.{f",{nl}    tmp.".join(columns)}
             );
             """,
         )

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/test/declarative/test_suites/__init__.py

@@ -6,6 +6,7 @@ Here we have base classes for a robust set of declarative connector test suites.
 
 from airbyte_cdk.test.declarative.test_suites.connector_base import (
     ConnectorTestScenario,
+    ConnectorTestSuiteBase,
     generate_tests,
 )
 from airbyte_cdk.test.declarative.test_suites.declarative_sources import (

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/test/declarative/test_suites/connector_base.py

@@ -4,34 +4,28 @@
 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"
+MANIFEST_YAML = "manifest.yaml"
 
 
 class JavaClass(str):
@@ -49,7 +43,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.
 
@@ -96,23 +90,21 @@ def generate_tests(metafunc) -> None:
 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.
-    """
+    @classmethod
+    def get_test_class_dir(cls) -> Path:
+        """Get the file path that contains the class."""
+        module = sys.modules[cls.__module__]
+        # Get the directory containing the test file
+        return Path(inspect.getfile(module)).parent
 
     connector: type[Connector] | Path | JavaClass | DockerImage | None = None
     """The connector class or path to the connector to test."""
 
-    working_dir: Path | None = None
-    """The root directory of the connector source code."""
-
     @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 @@ class ConnectorTestSuiteBase(abc.ABC):
         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,27 +138,43 @@ class ConnectorTestSuiteBase(abc.ABC):
             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
-    def acceptance_test_config_path(self) -> Path:
-        """Get the path to the acceptance test config file.
+    def get_connector_root_dir(cls) -> Path:
+        """Get the root directory of the connector."""
+        for parent in cls.get_test_class_dir().parents:
+            if (parent / MANIFEST_YAML).exists():
+                return parent
+            if (parent / ACCEPTANCE_TEST_CONFIG).exists():
+                return parent
+            if parent.name == "airbyte_cdk":
+                break
+        # If we reach here, we didn't find the manifest file in any parent directory
+        # Check if the manifest file exists in the current directory
+        for parent in Path.cwd().parents:
+            if (parent / MANIFEST_YAML).exists():
+                return parent
+            if (parent / ACCEPTANCE_TEST_CONFIG).exists():
+                return parent
+            if parent.name == "airbyte_cdk":
+                break
 
-        Check vwd and parent directories of cwd for the config file, and return the first one found.
-
-        Give up if the config file is not found in any parent directory.
-        """
-        current_dir = Path.cwd()
-        for parent_dir in current_dir.parents:
-            config_path = parent_dir / ACCEPTANCE_TEST_CONFIG
-            if config_path.exists():
-                return config_path
         raise FileNotFoundError(
-            f"Acceptance test config file not found in any parent directory from : {Path.cwd()}"
+            "Could not find connector root directory relative to "
+            f"'{str(cls.get_test_class_dir())}' or '{str(Path.cwd())}'."
         )
 
+    @classproperty
+    def acceptance_test_config_path(cls) -> Path:
+        """Get the path to the acceptance test config file."""
+        result = cls.get_connector_root_dir() / ACCEPTANCE_TEST_CONFIG
+        if result.exists():
+            return result
+
+        raise FileNotFoundError(f"Acceptance test config file not found at: {str(result)}")
+
     @classmethod
     def get_scenarios(
         cls,
@@ -193,10 +201,11 @@ class ConnectorTestSuiteBase(abc.ABC):
             for test in all_tests_config["acceptance_tests"][category]["tests"]
             if "iam_role" not in test["config_path"]
         ]
-        working_dir = cls.working_dir or Path()
+        connector_root = cls.get_connector_root_dir().absolute()
         for test in tests_scenarios:
             if test.config_path:
-                test.config_path = working_dir / test.config_path
+                test.config_path = connector_root / test.config_path
             if test.configured_catalog_path:
-                test.configured_catalog_path = working_dir / test.configured_catalog_path
+                test.configured_catalog_path = connector_root / test.configured_catalog_path
+
         return tests_scenarios

airbyte_cdk-6.45.1.post44.dev14417290834/airbyte_cdk/test/declarative/test_suites/declarative_sources.py

@@ -0,0 +1,74 @@
+import os
+from hashlib import md5
+from pathlib import Path
+from typing import Any, cast
+
+import yaml
+from boltons.typeutils import classproperty
+
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    ConcurrentDeclarativeSource,
+)
+from airbyte_cdk.test.declarative.models import ConnectorTestScenario
+from airbyte_cdk.test.declarative.test_suites.connector_base import MANIFEST_YAML
+from airbyte_cdk.test.declarative.test_suites.source_base import (
+    SourceTestSuiteBase,
+)
+from airbyte_cdk.test.declarative.utils.job_runner import IConnector
+
+
+def md5_checksum(file_path: Path) -> str:
+    with open(file_path, "rb") as file:
+        return md5(file.read()).hexdigest()
+
+
+class DeclarativeSourceTestSuite(SourceTestSuiteBase):
+    @classproperty
+    def manifest_yaml_path(cls) -> Path:
+        """Get the path to the manifest.yaml file."""
+        result = cls.get_connector_root_dir() / MANIFEST_YAML
+        if result.exists():
+            return result
+
+        raise FileNotFoundError(
+            f"Manifest YAML file not found at {result}. "
+            "Please ensure that the test suite is run in the correct directory.",
+        )
+
+    @classproperty
+    def components_py_path(cls) -> Path | None:
+        """Get the path to the components.py file."""
+        result = cls.get_connector_root_dir() / "components.py"
+        if result.exists():
+            return result
+
+        return None
+
+    @classmethod
+    def create_connector(
+        cls,
+        scenario: ConnectorTestScenario,
+    ) -> IConnector:
+        """Create a connector instance for the test suite."""
+        config: dict[str, Any] = scenario.get_config_dict()
+        # catalog = scenario.get_catalog()
+        # state = scenario.get_state()
+        # source_config = scenario.get_source_config()
+
+        manifest_dict = yaml.safe_load(cls.manifest_yaml_path.read_text())
+        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()
+            config["__injected_components_py_checksums"] = {
+                "md5": md5_checksum(cls.components_py_path),
+            }
+
+        return cast(
+            IConnector,
+            ConcurrentDeclarativeSource(
+                config=config,
+                catalog=None,
+                state=None,
+                source_config=manifest_dict,
+            ),
+        )

{airbyte_cdk-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/test/declarative/test_suites/source_base.py

@@ -2,9 +2,6 @@
 """Base class for source test suites."""
 
 from dataclasses import asdict
-from pathlib import Path
-
-import pytest
 
 from airbyte_cdk.models import (
     AirbyteMessage,
@@ -45,8 +42,10 @@ class SourceTestSuiteBase(ConnectorTestSuiteBase):
         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 +69,7 @@ class SourceTestSuiteBase(ConnectorTestSuiteBase):
                     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, union-attr]
             ]
         )
         result = run_test_job(
@@ -102,8 +101,8 @@ class SourceTestSuiteBase(ConnectorTestSuiteBase):
                         },
                         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-6.45.0.post24.dev14387467928 → airbyte_cdk-6.45.1.post44.dev14417290834}/airbyte_cdk/test/declarative/utils/job_runner.py

@@ -1,53 +1,62 @@
 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,
+        )
+        or True
+    ):  # TODO: Get a valid protocol check here
         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,14 +91,14 @@ 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 [arg-type]
         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],
             )
         )
 
@@ -106,10 +115,16 @@ def run_test_job(
             + "\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 +136,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
+                ],
             )
         )