fix test handling for expected errors

Author: aaronsteers

airbyte_cdk/test/declarative/test_suites/source_base.py

@@ -59,7 +59,10 @@ def test_basic_read(
             "discover",
             test_instance=instance,
         )
-        assert discover_result.catalog, "Expected a non-empty catalog."
+        if instance.expect_exception:
+            assert discover_result.errors, "Expected exception but got none."
+            return
+
         configured_catalog = ConfiguredAirbyteCatalog(
             streams=[
                 ConfiguredAirbyteStream(

airbyte_cdk/test/declarative/utils/job_runner.py

@@ -6,6 +6,9 @@
 import orjson
 
 from airbyte_cdk import Connector
+from airbyte_cdk.models import (
+    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
@@ -59,6 +62,7 @@ def run_test_job(
 
     catalog_path: Path | None = None
     if verb not in ["discover", "check"]:
+        # We need a catalog for read.
         if catalog:
             # Write the catalog to a temp json file and pass the path to the file as an argument.
             catalog_path = (
@@ -89,7 +93,36 @@ def run_test_job(
             )
         )
 
-    if test_instance.expect_exception and not result.errors:
-        raise AssertionError("Expected exception but got none.")  # noqa: TRY003
+    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 len(result.connection_status_messages) == 1, (
+            "Expected exactly one CONNECTION_STATUS message. Got "
+            f"{len(result.connection_status_messages)}:\n"
+            + "\n".join(result.connection_status_messages)
+        )
+        if test_instance.expect_exception:
+            assert result.connection_status_messages[0].connectionStatus.status == Status.FAILED, (
+                "Expected CONNECTION_STATUS message to be FAILED. Got: \n"
+                + "\n".join([str(result.connection_status_messages)])
+            )
+        return result
+
+    # For all other verbs, we assert check that an exception is raised (or not).
+    if test_instance.expect_exception:
+        if not result.errors:
+            raise AssertionError("Expected exception but got none.")
+
+        return result
+    if result.errors:
+        raise AssertionError(
+            "\n\n".join(
+                [str(err.trace.error).replace("\\n", "\n") for err in result.errors],
+            )
+        )
 
     return result

airbyte_cdk/test/entrypoint_wrapper.py

@@ -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 connection_status_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.CONNECTION_STATUS])
+
     @property
     def most_recent_state(self) -> Any:
         state_messages = self._get_message_by_types([Type.STATE])