improve tests

Author: aaronsteers

airbyte_cdk/test/standard_tests/_job_runner.py

@@ -148,8 +148,9 @@ def run_test_job(
 
         return result
 
-    assert not result.errors, (
-        f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
-    )
+    if not test_scenario.expect_exception:
+        assert not result.errors, (
+            f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
+        )
 
     return result

airbyte_cdk/test/standard_tests/connector_base.py

@@ -13,12 +13,10 @@
 from typing import cast
 
 import yaml
+from airbyte_protocol_dataclasses.models.airbyte_protocol import AirbyteConnectionStatus
 from boltons.typeutils import classproperty
 
-from airbyte_cdk.models import (
-    AirbyteMessage,
-    Type,
-)
+from airbyte_cdk.models import Status
 from airbyte_cdk.test import entrypoint_wrapper
 from airbyte_cdk.test.standard_tests._job_runner import IConnector, run_test_job
 from airbyte_cdk.test.standard_tests.models import (
@@ -117,12 +115,22 @@ def test_check(
             "check",
             test_scenario=scenario,
         )
-        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, (
-            f"Expected exactly one CONNECTION_STATUS message. Got: {result._messages}"
+        assert len(result.connection_status_messages) == 1, (
+            f"Expected exactly one CONNECTION_STATUS message. Got {len(result.connection_status_messages)}: \n"
+            + "\n".join([str(m) for m in result._messages])
+            + "\nErrors: "
+            + str(result.errors)
+            or "None"
+        )
+        conn_status = cast(
+            AirbyteConnectionStatus, result.connection_status_messages[0].connectionStatus
         )
+        if (
+            scenario.expect_exception
+            and conn_status.status == Status.SUCCEEDED
+            and not result.errors
+        ):
+            raise AssertionError(f"Expected error in `check` but got success.")
 
     @classmethod
     def get_connector_root_dir(cls) -> Path:

airbyte_cdk/test/standard_tests/source_base.py

@@ -151,6 +151,11 @@ def test_fail_read_with_bad_catalog(
         scenario: ConnectorTestScenario,
     ) -> None:
         """Standard test for `read` when passed a bad catalog file."""
+        # Recreate the scenario with the same config but set the status to "failed".
+        scenario = ConnectorTestScenario(
+            config_dict=scenario.get_config_dict(empty_if_missing=False),
+            status="failed",
+        )
         invalid_configured_catalog = ConfiguredAirbyteCatalog(
             streams=[
                 # Create ConfiguredAirbyteStream which is deliberately invalid
@@ -171,7 +176,6 @@ def test_fail_read_with_bad_catalog(
             ]
         )
         # Set expected status to "failed" to ensure the test fails if the connector.
-        scenario.status = "failed"
         result: entrypoint_wrapper.EntrypointOutput = run_test_job(
             self.create_connector(scenario),
             "read",