Skip to content
4 changes: 4 additions & 0 deletions airbyte_cdk/test/entrypoint_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
24 changes: 15 additions & 9 deletions airbyte_cdk/test/standard_tests/_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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"
)
Expand All @@ -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:
Expand All @@ -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=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)
)
Expand All @@ -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"
Expand All @@ -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.")

Expand Down
35 changes: 21 additions & 14 deletions airbyte_cdk/test/standard_tests/connector_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -147,28 +147,35 @@ 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"]:
return []
if "tests" not in all_tests_config["acceptance_tests"][category]:
raise ValueError(f"No tests found for category {category}")

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
10 changes: 7 additions & 3 deletions airbyte_cdk/test/standard_tests/declarative_sources.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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()
Expand Down
20 changes: 20 additions & 0 deletions airbyte_cdk/test/standard_tests/source_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading