tri-state ExpectedOutcome

Author: aaronsteers

airbyte_cdk/test/entrypoint_wrapper.py

@@ -24,6 +24,7 @@
 from typing import Any, List, Mapping, Optional, Union
 
 import orjson
+from langsmith import expect
 from pydantic import ValidationError as V2ValidationError
 from serpyco_rs import SchemaValidationError
 
@@ -44,6 +45,7 @@
     Type,
 )
 from airbyte_cdk.sources import Source
+from airbyte_cdk.test.standard_tests.models.scenario import ExpectedOutcome
 
 
 class EntrypointOutput:
@@ -157,7 +159,9 @@ def is_not_in_logs(self, pattern: str) -> bool:
 
 
 def _run_command(
-    source: Source, args: List[str], expecting_exception: bool = False
+    source: Source,
+    args: list[str],
+    expected_outcome: ExpectedOutcome,
 ) -> EntrypointOutput:
     log_capture_buffer = StringIO()
     stream_handler = logging.StreamHandler(log_capture_buffer)
@@ -175,35 +179,39 @@ def _run_command(
         for message in source_entrypoint.run(parsed_args):
             messages.append(message)
     except Exception as exception:
-        if not expecting_exception:
+        if expected_outcome.expect_success():
             print("Printing unexpected error from entrypoint_wrapper")
             print("".join(traceback.format_exception(None, exception, exception.__traceback__)))
+
         uncaught_exception = exception
 
     captured_logs = log_capture_buffer.getvalue().split("\n")[:-1]
 
     parent_logger.removeHandler(stream_handler)
 
-    return EntrypointOutput(messages + captured_logs, uncaught_exception)
+    return EntrypointOutput(messages + captured_logs, uncaught_exception=uncaught_exception)
 
 
 def discover(
     source: Source,
     config: Mapping[str, Any],
-    expecting_exception: bool = False,
+    *,
+    expected_outcome: ExpectedOutcome = ExpectedOutcome.EXPECT_SUCCESS,
 ) -> EntrypointOutput:
     """
     config must be json serializable
-    :param expecting_exception: By default if there is an uncaught exception, the exception will be printed out. If this is expected, please
-        provide expecting_exception=True so that the test output logs are cleaner
+    :param expected_outcome: By default if there is an uncaught exception, the exception will be printed out. If this is expected, please
+        provide `expected_outcome=ExpectedOutcome.EXPECT_FAILURE` so that the test output logs are cleaner
     """
 
     with tempfile.TemporaryDirectory() as tmp_directory:
         tmp_directory_path = Path(tmp_directory)
         config_file = make_file(tmp_directory_path / "config.json", config)
 
         return _run_command(
-            source, ["discover", "--config", config_file, "--debug"], expecting_exception
+            source,
+            ["discover", "--config", config_file, "--debug"],
+            expected_outcome=expected_outcome,
         )
 
 
@@ -212,13 +220,14 @@ def read(
     config: Mapping[str, Any],
     catalog: ConfiguredAirbyteCatalog,
     state: Optional[List[AirbyteStateMessage]] = None,
-    expecting_exception: bool = False,
+    *,
+    expected_outcome: ExpectedOutcome = ExpectedOutcome.EXPECT_SUCCESS,
 ) -> EntrypointOutput:
     """
     config and state must be json serializable
 
-    :param expecting_exception: By default if there is an uncaught exception, the exception will be printed out. If this is expected, please
-        provide expecting_exception=True so that the test output logs are cleaner
+    :param expected_outcome: By default if there is an uncaught exception, the exception will be printed out. If this is expected, please
+        provide `expected_outcome=ExpectedOutcome.EXPECT_FAILURE` so that the test output logs are cleaner.
     """
     with tempfile.TemporaryDirectory() as tmp_directory:
         tmp_directory_path = Path(tmp_directory)
@@ -245,7 +254,7 @@ def read(
                 ]
             )
 
-        return _run_command(source, args, expecting_exception)
+        return _run_command(source, args, expected_outcome=expected_outcome)
 
 
 def make_file(

airbyte_cdk/test/standard_tests/_job_runner.py

@@ -122,9 +122,9 @@ def run_test_job(
     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,
+        expected_outcome=test_scenario.expected_outcome,
     )
-    if result.errors and not test_scenario.expect_exception:
+    if result.errors and test_scenario.expected_outcome.expect_success():
         raise AssertionError(
             f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
         )
@@ -139,7 +139,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 test_scenario.expected_outcome.expect_exception():
             conn_status = result.connection_status_messages[0].connectionStatus
             assert conn_status, (
                 "Expected CONNECTION_STATUS message to be present. Got: \n"
@@ -153,14 +153,15 @@ 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 test_scenario.expected_outcome.expect_exception():
         if not result.errors:
             raise AssertionError("Expected exception but got none.")
 
         return result
 
-    assert not result.errors, (
-        f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
-    )
+    if test_scenario.expected_outcome.expect_success():
+        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/models/scenario.py

@@ -9,13 +9,47 @@
 
 from __future__ import annotations
 
-from pathlib import Path
+from enum import Enum, auto
+from pathlib import Path  # noqa: TC003  # Pydantic needs this (don't move to 'if typing' block)
 from typing import Any, Literal, cast
 
 import yaml
 from pydantic import BaseModel, ConfigDict
 
 
+class ExpectedOutcome(Enum):
+    """Enum to represent the expected outcome of a test scenario.
+
+    Class supports comparisons to a boolean or None.
+    """
+
+    EXPECT_EXCEPTION = auto()
+    EXPECT_SUCCESS = auto()
+    ALLOW_ANY = auto()
+
+    @classmethod
+    def from_status_str(cls, status: str | None) -> ExpectedOutcome:
+        """Convert a status string to an ExpectedOutcome."""
+        if status is None:
+            return ExpectedOutcome.ALLOW_ANY
+
+        try:
+            return {
+                "succeed": ExpectedOutcome.EXPECT_SUCCESS,
+                "failed": ExpectedOutcome.EXPECT_EXCEPTION,
+            }[status]
+        except KeyError as ex:
+            raise ValueError(f"Invalid status '{status}'. Expected 'succeed' or 'failed'.") from ex
+
+    def expect_exception(self) -> bool:
+        """Return whether the expectation is that an exception should be raised."""
+        return self == ExpectedOutcome.EXPECT_EXCEPTION
+
+    def expect_success(self) -> bool:
+        """Return whether the expectation is that the test should succeed without exceptions."""
+        return self == ExpectedOutcome.EXPECT_SUCCESS
+
+
 class ConnectorTestScenario(BaseModel):
     """Acceptance test scenario, as a Pydantic model.
 
@@ -82,8 +116,13 @@ def get_config_dict(
         raise ValueError("No config dictionary or path provided.")
 
     @property
-    def expect_exception(self) -> bool:
-        return self.status and self.status == "failed" or False
+    def expected_outcome(self) -> ExpectedOutcome:
+        """Whether the test scenario expects an exception to be raised.
+
+        Returns True if the scenario expects an exception, False if it does not,
+        and None if there is no set expectation.
+        """
+        return ExpectedOutcome.from_status_str(self.status)
 
     @property
     def instance_name(self) -> str:
@@ -97,15 +136,11 @@ def __str__(self) -> str:
 
         return f"'{hash(self)}' Test Scenario"
 
-    def without_expecting_failure(self) -> ConnectorTestScenario:
-        """Return a copy of the scenario that does not expect failure.
+    def without_expected_outcome(self) -> ConnectorTestScenario:
+        """Return a copy of the scenario that does not expect failure or success.
 
-        This is useful when you need to run multiple steps and you
-        want to defer failure expectation for one or more steps.
+        This is useful when running multiple steps, to defer the expectations to a later step.
         """
-        if self.status != "failed":
-            return self
-
         return ConnectorTestScenario(
             **self.model_dump(exclude={"status"}),
         )
@@ -122,3 +157,16 @@ def with_expecting_failure(self) -> ConnectorTestScenario:
             **self.model_dump(exclude={"status"}),
             status="failed",
         )
+
+    def with_expecting_success(self) -> ConnectorTestScenario:
+        """Return a copy of the scenario that expects success.
+
+        This is useful when deriving new scenarios from existing ones.
+        """
+        if self.status == "succeed":
+            return self
+
+        return ConnectorTestScenario(
+            **self.model_dump(exclude={"status"}),
+            status="succeed",
+        )

airbyte_cdk/test/standard_tests/source_base.py

@@ -106,9 +106,9 @@ def test_basic_read(
             self.create_connector(scenario),
             "discover",
             connector_root=self.get_connector_root_dir(),
-            test_scenario=scenario.without_expecting_failure(),
+            test_scenario=scenario.without_expected_outcome(),
         )
-        if scenario.expect_exception and discover_result.errors:
+        if scenario.expected_outcome.expect_exception() and discover_result.errors:
             # Failed as expected; we're done.
             return
 

airbyte_cdk/test/utils/reading.py

@@ -6,6 +6,7 @@
 from airbyte_cdk.models import AirbyteStateMessage, ConfiguredAirbyteCatalog, SyncMode
 from airbyte_cdk.test.catalog_builder import CatalogBuilder
 from airbyte_cdk.test.entrypoint_wrapper import EntrypointOutput, read
+from airbyte_cdk.test.standard_tests.models.scenario import ExpectedOutcome
 
 
 def catalog(stream_name: str, sync_mode: SyncMode) -> ConfiguredAirbyteCatalog:
@@ -19,8 +20,8 @@ def read_records(
     stream_name: str,
     sync_mode: SyncMode,
     state: Optional[List[AirbyteStateMessage]] = None,
-    expecting_exception: bool = False,
+    expected_outcome: ExpectedOutcome = ExpectedOutcome.EXPECT_SUCCESS,
 ) -> EntrypointOutput:
     """Read records from a stream."""
     _catalog = catalog(stream_name, sync_mode)
-    return read(source, config, _catalog, state, expecting_exception)
+    return read(source, config, _catalog, state, expected_outcome=expected_outcome)

unit_tests/sources/declarative/file/test_file_stream.py

@@ -17,6 +17,7 @@
 from airbyte_cdk.test.entrypoint_wrapper import read as entrypoint_read
 from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse
 from airbyte_cdk.test.mock_http.response_builder import find_binary_response, find_template
+from airbyte_cdk.test.standard_tests.models.scenario import ExpectedOutcome
 from airbyte_cdk.test.state_builder import StateBuilder
 
 
@@ -53,8 +54,9 @@ def read(
     config_builder: ConfigBuilder,
     catalog: ConfiguredAirbyteCatalog,
     state_builder: Optional[StateBuilder] = None,
-    expecting_exception: bool = False,
     yaml_file: Optional[str] = None,
+    *,
+    expected_outcome: ExpectedOutcome = ExpectedOutcome.EXPECT_SUCCESS,
 ) -> EntrypointOutput:
     config = config_builder.build()
     state = state_builder.build() if state_builder else StateBuilder().build()
@@ -63,14 +65,19 @@ def read(
         config,
         catalog,
         state,
-        expecting_exception,
+        expected_outcome=expected_outcome,
     )
 
 
-def discover(config_builder: ConfigBuilder, expecting_exception: bool = False) -> EntrypointOutput:
+def discover(
+    config_builder: ConfigBuilder,
+    expected_outcome: ExpectedOutcome = ExpectedOutcome.EXPECT_SUCCESS,
+) -> EntrypointOutput:
     config = config_builder.build()
     return entrypoint_discover(
-        _source(CatalogBuilder().build(), config), config, expecting_exception
+        _source(CatalogBuilder().build(), config),
+        config,
+        expected_outcome=expected_outcome,
     )
 
 

unit_tests/sources/mock_server_tests/test_resumable_full_refresh.py

@@ -27,6 +27,7 @@
     create_record_builder,
     create_response_builder,
 )
+from airbyte_cdk.test.standard_tests.models.scenario import ExpectedOutcome
 from airbyte_cdk.test.state_builder import StateBuilder
 from unit_tests.sources.mock_server_tests.mock_source_fixture import SourceFixture
 from unit_tests.sources.mock_server_tests.test_helpers import (
@@ -344,7 +345,7 @@ def test_resumable_full_refresh_failure(self, http_mocker):
             source,
             config=config,
             catalog=_create_catalog([("justice_songs", SyncMode.full_refresh, {})]),
-            expecting_exception=True,
+            expected_outcome=ExpectedOutcome.EXPECT_EXCEPTION,
         )
 
         status_messages = actual_messages.get_stream_statuses("justice_songs")

unit_tests/test/test_entrypoint_wrapper.py

@@ -32,6 +32,7 @@
 )
 from airbyte_cdk.sources.abstract_source import AbstractSource
 from airbyte_cdk.test.entrypoint_wrapper import EntrypointOutput, discover, read
+from airbyte_cdk.test.standard_tests.models.scenario import ExpectedOutcome
 from airbyte_cdk.test.state_builder import StateBuilder
 
 
@@ -229,7 +230,7 @@ def test_given_unexpected_exception_when_discover_then_print(self, entrypoint, p
     @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
     def test_given_expected_exception_when_discover_then_do_not_print(self, entrypoint, print_mock):
         entrypoint.return_value.run.side_effect = ValueError("This error should not be printed")
-        discover(self._a_source, _A_CONFIG, expecting_exception=True)
+        discover(self._a_source, _A_CONFIG, expected_outcome=ExpectedOutcome.EXPECT_EXCEPTION)
         assert print_mock.call_count == 0
 
     @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
@@ -380,7 +381,13 @@ def test_given_unexpected_exception_when_read_then_print(self, entrypoint, print
     @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
     def test_given_expected_exception_when_read_then_do_not_print(self, entrypoint, print_mock):
         entrypoint.return_value.run.side_effect = ValueError("This error should not be printed")
-        read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE, expecting_exception=True)
+        read(
+            self._a_source,
+            _A_CONFIG,
+            _A_CATALOG,
+            _A_STATE,
+            expected_outcome=ExpectedOutcome.EXPECT_EXCEPTION,
+        )
         assert print_mock.call_count == 0
 
     @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")