airbyte-cdk 0.55.3__py3-none-any.whl → 0.55.5__py3-none-any.whl

airbyte_cdk/entrypoint.py

@@ -106,6 +106,9 @@ class AirbyteEntrypoint(object):
                     raw_config = self.source.read_config(parsed_args.config)
                     config = self.source.configure(raw_config, temp_dir)
 
+                    yield from [
+                        self.airbyte_message_to_string(queued_message) for queued_message in self._emit_queued_messages(self.source)
+                    ]
                     if cmd == "check":
                         yield from map(AirbyteEntrypoint.airbyte_message_to_string, self.check(source_spec, config))
                     elif cmd == "discover":

airbyte_cdk/sources/abstract_source.py

@@ -19,7 +19,7 @@ from airbyte_cdk.models import (
 )
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
-from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
 from airbyte_cdk.sources.source import Source
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.core import StreamData
@@ -31,6 +31,8 @@ from airbyte_cdk.utils.event_timing import create_timer
 from airbyte_cdk.utils.stream_status_utils import as_airbyte_message as stream_status_as_airbyte_message
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
+_default_message_repository = InMemoryMessageRepository()
+
 
 class AbstractSource(Source, ABC):
     """
@@ -269,4 +271,4 @@ class AbstractSource(Source, ABC):
 
     @property
     def message_repository(self) -> Union[None, MessageRepository]:
-        return None
+        return _default_message_repository

airbyte_cdk/sources/file_based/config/file_based_stream_config.py

@@ -46,7 +46,9 @@ class FileBasedStreamConfig(BaseModel):
         description="The schema that will be used to validate records extracted from the file. This will override the stream schema that is auto-detected from incoming files.",
     )
     primary_key: Optional[str] = Field(
-        title="Primary Key", description="The column or columns (for a composite key) that serves as the unique identifier of a record."
+        title="Primary Key",
+        description="The column or columns (for a composite key) that serves as the unique identifier of a record.",
+        airbyte_hidden=True,  # Users can create/modify primary keys in the connection configuration so we shouldn't duplicate it here.
     )
     days_to_sync_if_history_is_full: int = Field(
         title="Days To Sync If History Is Full",

airbyte_cdk/test/entrypoint_wrapper.py

@@ -0,0 +1,116 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+
+"""
+The AirbyteEntrypoint is important because it is a service layer that orchestrate how we execute commands from the
+[common interface](https://docs.airbyte.com/understanding-airbyte/airbyte-protocol#common-interface) through the source Python
+implementation. There is some logic about which message we send to the platform and when which is relevant for integration testing. Other
+than that, there are integrations point that are annoying to integrate with using Python code:
+* Sources communicate with the platform using stdout. The implication is that the source could just print every message instead of
+    returning things to source.<method> or to using the message repository. WARNING: As part of integration testing, we will not support
+    messages that are simply printed. The reason is that capturing stdout relies on overriding sys.stdout (see
+    https://docs.python.org/3/library/contextlib.html#contextlib.redirect_stdout) which clashes with how pytest captures logs and brings
+    considerations for multithreaded applications. If code you work with uses `print` statements, please migrate to
+    source.message_repository to emit those messages
+* The entrypoint interface relies on file being written on the file system
+"""
+
+import json
+import logging
+import tempfile
+from io import StringIO
+from pathlib import Path
+from typing import Any, List, Mapping, Optional, Union
+
+from airbyte_cdk.entrypoint import AirbyteEntrypoint
+from airbyte_cdk.logger import AirbyteLogFormatter
+from airbyte_cdk.sources import Source
+from airbyte_protocol.models import AirbyteLogMessage, AirbyteMessage, ConfiguredAirbyteCatalog, Level, TraceType, Type
+from pydantic.error_wrappers import ValidationError
+
+
+class EntrypointOutput:
+    def __init__(self, messages: List[str]):
+        try:
+            self._messages = [self._parse_message(message) for message in messages]
+        except ValidationError as exception:
+            raise ValueError("All messages are expected to be AirbyteMessage") from exception
+
+    @staticmethod
+    def _parse_message(message: str) -> AirbyteMessage:
+        try:
+            return AirbyteMessage.parse_obj(json.loads(message))
+        except (json.JSONDecodeError, ValidationError):
+            # The platform assumes that logs that are not of AirbyteMessage format are log messages
+            return AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message=message))
+
+    @property
+    def records_and_state_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.RECORD, Type.STATE])
+
+    @property
+    def records(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.RECORD])
+
+    @property
+    def state_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.STATE])
+
+    @property
+    def logs(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.LOG])
+
+    @property
+    def trace_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.TRACE])
+
+    @property
+    def analytics_messages(self) -> List[AirbyteMessage]:
+        return [message for message in self._get_message_by_types([Type.TRACE]) if message.trace.type == TraceType.ANALYTICS]
+
+    def _get_message_by_types(self, message_types: List[Type]) -> List[AirbyteMessage]:
+        return [message for message in self._messages if message.type in message_types]
+
+
+def read(source: Source, config: Mapping[str, Any], catalog: ConfiguredAirbyteCatalog, state: Optional[Any] = None) -> EntrypointOutput:
+    """
+    config and state must be json serializable
+    """
+    log_capture_buffer = StringIO()
+    stream_handler = logging.StreamHandler(log_capture_buffer)
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(AirbyteLogFormatter())
+    parent_logger = logging.getLogger("")
+    parent_logger.addHandler(stream_handler)
+
+    with tempfile.TemporaryDirectory() as tmp_directory:
+        tmp_directory_path = Path(tmp_directory)
+        args = [
+            "read",
+            "--config",
+            make_file(tmp_directory_path / "config.json", config),
+            "--catalog",
+            make_file(tmp_directory_path / "catalog.json", catalog.json()),
+        ]
+        if state:
+            args.extend(
+                [
+                    "--state",
+                    make_file(tmp_directory_path / "state.json", state),
+                ]
+            )
+        source_entrypoint = AirbyteEntrypoint(source)
+        parsed_args = source_entrypoint.parse_args(args)
+        messages = list(source_entrypoint.run(parsed_args))
+        captured_logs = log_capture_buffer.getvalue().split("\n")[:-1]
+
+        parent_logger.removeHandler(stream_handler)
+
+        return EntrypointOutput(messages + captured_logs)
+
+
+def make_file(path: Path, file_contents: Optional[Union[str, Mapping[str, Any], List[Mapping[str, Any]]]]) -> str:
+    if isinstance(file_contents, str):
+        path.write_text(file_contents)
+    else:
+        path.write_text(json.dumps(file_contents))
+    return str(path)

{airbyte_cdk-0.55.3.dist-info → airbyte_cdk-0.55.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 0.55.3
+Version: 0.55.5
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://github.com/airbytehq/airbyte
 Author: Airbyte

{airbyte_cdk-0.55.3.dist-info → airbyte_cdk-0.55.5.dist-info}/RECORD

@@ -1,7 +1,7 @@
 airbyte_cdk/__init__.py,sha256=OBQWv5rF_QTRpOiP6J8J8oTU-GGrfi18i1PRFpahKks,262
 airbyte_cdk/config_observation.py,sha256=3kjxv8xTwCnub2_fTWnMPRx0E7vly1BUeyXOSK15Ql4,3610
 airbyte_cdk/connector.py,sha256=LtTAmBFV1LBUz_fOEbQ_EvBhyUsz8AGOlDsvK8QOOo0,4396
-airbyte_cdk/entrypoint.py,sha256=upN2KOPYlyalns49ZmZVuAnEMVDxXb33lv5NUEnxUiY,12992
+airbyte_cdk/entrypoint.py,sha256=uX3MawH1qukzxFjdR1AFynG0l5vLof9X40m_AYTrP_8,13180
 airbyte_cdk/exception_handler.py,sha256=CwkiPdZ1WMOr3CBkvKFyHiyLerXGRqBrVlB4p0OImGI,1125
 airbyte_cdk/logger.py,sha256=4Mi2MEQi1uh59BP9Dxw_UEbZuxaJewqK_jvEU2b10nk,3985
 airbyte_cdk/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -24,7 +24,7 @@ airbyte_cdk/models/__init__.py,sha256=Kg8YHBqUsNWHlAw-u3ZGdG4dxLh7qBlHhqMRfamNCR
 airbyte_cdk/models/airbyte_protocol.py,sha256=DoJvnmGM3xMAZFTwA6_RGMiKSFqfE3ib_Ru0KJ65Ag4,100
 airbyte_cdk/models/well_known_types.py,sha256=KKfNbow2gdLoC1Z4hcXy_JR8m_acsB2ol7gQuEgjobw,117
 airbyte_cdk/sources/__init__.py,sha256=Ov7Uf03KPSZUmMZqZfUAK3tQwsdKjDQUDvTb-H0JyfA,1141
-airbyte_cdk/sources/abstract_source.py,sha256=_cUJPfxNrYIq7upQ17wQeIchC7jZYWDB-6XfDNA-yCM,12934
+airbyte_cdk/sources/abstract_source.py,sha256=tlXXwCRBisbOu7pA2NYtiRolhHSCaKlXcET6dny5V9o,13043
 airbyte_cdk/sources/config.py,sha256=PYsY7y2u3EUwxLiEb96JnuKwH_E8CuxKggsRO2ZPSRc,856
 airbyte_cdk/sources/connector_state_manager.py,sha256=wsmUgII398MazCTKxwLBLzeiU6Z-tMTrKX882EEy-YE,10904
 airbyte_cdk/sources/http_config.py,sha256=OBZeuyFilm6NlDlBhFQvHhTWabEvZww6OHDIlZujIS0,730
@@ -164,7 +164,7 @@ airbyte_cdk/sources/file_based/config/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCe
 airbyte_cdk/sources/file_based/config/abstract_file_based_spec.py,sha256=WrV4sKtJoZ1dK31HK7NdBKlnYHkmu6NqjmEpkVqJ6tQ,4582
 airbyte_cdk/sources/file_based/config/avro_format.py,sha256=lQSEq5JZY0M5y9mW93R4EjrIb8brYXUgrXCY-6EMHww,711
 airbyte_cdk/sources/file_based/config/csv_format.py,sha256=L3JEgb91yrCob1oYrGl0088QEWblkOsRfDmMfWRQ0bg,7482
-airbyte_cdk/sources/file_based/config/file_based_stream_config.py,sha256=zzD59S8e9aeYfVU5Dh_i7dpFe2qi_thQyLpQYNQcnm0,4008
+airbyte_cdk/sources/file_based/config/file_based_stream_config.py,sha256=lE1_Uzcfoly63xIoNigRFT0urOKjf0FQ7fLMKH2AX-s,4153
 airbyte_cdk/sources/file_based/config/jsonl_format.py,sha256=fAPzZnoghGgHjaDvx6Qo68C8j54mBxo1NTdpwSI0VZo,374
 airbyte_cdk/sources/file_based/config/parquet_format.py,sha256=8GTDTQyvS7pWLVG0LWirHVE1snHd0Au5R4Ym33-ezEg,736
 airbyte_cdk/sources/file_based/config/unstructured_format.py,sha256=8yc0TMhlf1bcJc34IXzYkYHQ5HpGN4rt1f3zKSiCeYk,934
@@ -237,6 +237,7 @@ airbyte_cdk/sources/utils/slice_logger.py,sha256=YeWSoZeOsQp9oZK7mick2J8KFdiY726
 airbyte_cdk/sources/utils/transform.py,sha256=4GYmO6bq33HF-a1in0dKQKqUOYI1bWItyuYF875bSQg,9493
 airbyte_cdk/sources/utils/types.py,sha256=41ZQR681t5TUnOScij58d088sb99klH_ZENFcaYro_g,175
 airbyte_cdk/test/__init__.py,sha256=f_XdkOg4_63QT2k3BbKY34209lppwgw-svzfZstQEq4,199
+airbyte_cdk/test/entrypoint_wrapper.py,sha256=Bu67gvYwlnejmRc4fifkn4wKNOof9TuNfinYRdUUWVE,5155
 airbyte_cdk/test/http/__init__.py,sha256=Gh2u6y10KXfvYqhC3Mm811U-b9FaoQMBxSW4hT2JJ6I,302
 airbyte_cdk/test/http/matcher.py,sha256=zkcnnteguIYLI8nrZ1egttw94OfQPgY1LUsPjjytYgY,1117
 airbyte_cdk/test/http/mocker.py,sha256=pKEEHnf_OsOv-czx1jyyftfvCy4fDgtl6k7SweytOrU,4472
@@ -356,9 +357,9 @@ unit_tests/sources/declarative/stream_slicers/test_cartesian_product_stream_slic
 unit_tests/sources/file_based/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unit_tests/sources/file_based/helpers.py,sha256=MZTwaWtX0a6TPbFcUMP-EgqBunK2wpoElgApCEE1bN4,2659
 unit_tests/sources/file_based/in_memory_files_source.py,sha256=r2yD6-_ABXG7_PIyTq4ACN21sHyg3g-Hd9dIgxfDQUk,8235
-unit_tests/sources/file_based/test_file_based_scenarios.py,sha256=9xVFaFFHjnzZziVmoVmLTULdxANt_zSrwVgANAVytl4,11564
+unit_tests/sources/file_based/test_file_based_scenarios.py,sha256=CMlhlccQzVq65cqu8V_Kpo1Eafumv-2OPCRwWNBnbGk,11418
 unit_tests/sources/file_based/test_file_based_stream_reader.py,sha256=P6yTp7tbPfREzi5SXg4SSSql5nxiRV571YdOmwb_SzY,9219
-unit_tests/sources/file_based/test_scenarios.py,sha256=2-9pqnfva3RDRyODy0xcK6mxrP_mHH5vLrmBhqgZO8o,9703
+unit_tests/sources/file_based/test_scenarios.py,sha256=4xRRBZhqe__QWP8Fd5oXCHDI-0qKjrhypPAWGvmWwDM,7888
 unit_tests/sources/file_based/test_schema_helpers.py,sha256=IYIDdLRK41RkSG_ZW2cagAt9krV4QLbkzu6r7vPx9Js,12047
 unit_tests/sources/file_based/availability_strategy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unit_tests/sources/file_based/availability_strategy/test_default_file_based_availability_strategy.py,sha256=V9knz_3FHzk8k6R1JaiZgZjHHFNRa8pmH2PGqhRJ2UA,4432
@@ -377,12 +378,12 @@ unit_tests/sources/file_based/file_types/test_unstructured_parser.py,sha256=W6jb
 unit_tests/sources/file_based/scenarios/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unit_tests/sources/file_based/scenarios/avro_scenarios.py,sha256=oeQUmCV7d2aTShreYc-PvVb4cWqLSsVwHfg-lcKjzPs,30554
 unit_tests/sources/file_based/scenarios/check_scenarios.py,sha256=0xkt21ASTnTAMP0RYJEsF3yMGsNN7wWOoG_tmzL9PYw,6750
-unit_tests/sources/file_based/scenarios/csv_scenarios.py,sha256=PAvWvqUbmGyTts8xWSSx1lrMeXys1hOJ5b2jzfZaqNk,109990
+unit_tests/sources/file_based/scenarios/csv_scenarios.py,sha256=1u2oZse5VDVJVZBM4ZeLzP-6E7BDjhYCxHvWhQGFWBY,110050
 unit_tests/sources/file_based/scenarios/file_based_source_builder.py,sha256=wgb7l5VohcEvZT82ZpJcjINSrjuJtzJZS4zuZjdKpJ4,3874
 unit_tests/sources/file_based/scenarios/incremental_scenarios.py,sha256=B7YE2IbvgTH_v7DYQEuv7yn2IG15aKUvJ_7dA4d3Cg4,69413
 unit_tests/sources/file_based/scenarios/jsonl_scenarios.py,sha256=LsOf-tpjWNuwskPcgAMhMpQQ3iaHaD3PjPmt2M2zSzo,31839
 unit_tests/sources/file_based/scenarios/parquet_scenarios.py,sha256=MGgLCqkTJb8uNEwYZY3zbVVDZRSBKSmf2s8VMuYse_I,26549
-unit_tests/sources/file_based/scenarios/scenario_builder.py,sha256=feSSViayuoxTquoRhMUg4Lcui7dtwWHQ1Fe5y9igWSo,8728
+unit_tests/sources/file_based/scenarios/scenario_builder.py,sha256=GeuAOg7eC6Wgx648-OoAcaQh19AdFLgeLTpN2vmyTaw,8894
 unit_tests/sources/file_based/scenarios/unstructured_scenarios.py,sha256=3jeDxyLh6LgwK0wMhU884fqSXG47H3AWvIQDD15jO6c,64973
 unit_tests/sources/file_based/scenarios/user_input_schema_scenarios.py,sha256=7CxIaqZxAGSPs4AtcKZ9FLVVYQPsS__uXi9wnQMKn3U,28322
 unit_tests/sources/file_based/scenarios/validation_policy_scenarios.py,sha256=Try0knJN5wfoGNO38QGoLGIcqSceSAQsUWO42CusNYI,33005
@@ -411,7 +412,7 @@ unit_tests/sources/streams/concurrent/scenarios/__init__.py,sha256=4Hw-PX1-VgESL
 unit_tests/sources/streams/concurrent/scenarios/incremental_scenarios.py,sha256=x77AQf8_O4dQ2aF1o800CzI0hOEyU8ayxoNdSOvxkhM,10495
 unit_tests/sources/streams/concurrent/scenarios/stream_facade_builder.py,sha256=OD_9R5fHt5Nf7hH8m28-UDoZJkY8iUBJLI73kd-u2BE,5794
 unit_tests/sources/streams/concurrent/scenarios/stream_facade_scenarios.py,sha256=v0yP5MRGYJAb9bp2yXnp5yUmYKJ6aAKjHcNHigL_ONY,13981
-unit_tests/sources/streams/concurrent/scenarios/test_concurrent_scenarios.py,sha256=sQpvIJa5-Iv03KZfC2sP2zB8XSPCZAjLpUMpNBOA-xM,3897
+unit_tests/sources/streams/concurrent/scenarios/test_concurrent_scenarios.py,sha256=Z_4-ClsxBupmN7Pbl8lF9bkSA9wnjLtrgA9WR_8VRi8,3757
 unit_tests/sources/streams/concurrent/scenarios/thread_based_concurrent_stream_scenarios.py,sha256=KqCLsXB_9rV4hNdSPrNynK3G-UIsipqsZT6X0Z-iM5E,13175
 unit_tests/sources/streams/concurrent/scenarios/thread_based_concurrent_stream_source_builder.py,sha256=aMtEOpCkxH-v2BBOYj4xABzPKcDYh_jieGfaIp4hy9w,5727
 unit_tests/sources/streams/concurrent/scenarios/utils.py,sha256=Pl1F4asW8AvV6bV5W3Qg21GiLqfdMT_rOt1CsFA0aVM,1953
@@ -436,8 +437,8 @@ unit_tests/utils/test_schema_inferrer.py,sha256=Z2jHBZ540wnYkylIdV_2xr75Vtwlxuyg
 unit_tests/utils/test_secret_utils.py,sha256=XKe0f1RHYii8iwE6ATmBr5JGDI1pzzrnZUGdUSMJQP4,4886
 unit_tests/utils/test_stream_status_utils.py,sha256=Xr8MZ2HWgTVIyMbywDvuYkRaUF4RZLQOT8-JjvcfR24,2970
 unit_tests/utils/test_traced_exception.py,sha256=bDFP5zMBizFenz6V2WvEZTRCKGB5ijh3DBezjbfoYIs,4198
-airbyte_cdk-0.55.3.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
-airbyte_cdk-0.55.3.dist-info/METADATA,sha256=NMmUmx3R4MAMSEd1mEkbL2kl7wiOK2Yta6L6AnvVxMU,11983
-airbyte_cdk-0.55.3.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-airbyte_cdk-0.55.3.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
-airbyte_cdk-0.55.3.dist-info/RECORD,,
+airbyte_cdk-0.55.5.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
+airbyte_cdk-0.55.5.dist-info/METADATA,sha256=UGBscoQLxD10YK6vK2sziE_7IkU_JY03BSihYZcEJv0,11983
+airbyte_cdk-0.55.5.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+airbyte_cdk-0.55.5.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
+airbyte_cdk-0.55.5.dist-info/RECORD,,

unit_tests/sources/file_based/scenarios/csv_scenarios.py

@@ -100,6 +100,7 @@ single_csv_scenario: TestScenario[InMemoryFilesSource] = (
                                     "title": "Primary Key",
                                     "description": "The column or columns (for a composite key) that serves as the unique identifier of a record.",
                                     "type": "string",
+                                    "airbyte_hidden": True,
                                 },
                                 "days_to_sync_if_history_is_full": {
                                     "title": "Days To Sync If History Is Full",

unit_tests/sources/file_based/scenarios/scenario_builder.py

@@ -81,9 +81,13 @@ class TestScenario(Generic[SourceType]):
         for stream in self.source.streams(self.config):
             catalog["streams"].append(
                 {
-                    "stream": stream.name,
+                    "stream": {
+                        "name": stream.name,
+                        "json_schema": {},
+                        "supported_sync_modes": [sync_mode.value],
+                    },
                     "sync_mode": sync_mode.value,
-                    "destination_sync_mode": "append",
+                    "destination_sync_mode": "append"
                 }
             )
 

unit_tests/sources/file_based/test_file_based_scenarios.py

@@ -8,7 +8,6 @@ import pytest
 from _pytest.capture import CaptureFixture
 from airbyte_cdk.sources.abstract_source import AbstractSource
 from freezegun import freeze_time
-from pytest import LogCaptureFixture
 from unit_tests.sources.file_based.scenarios.avro_scenarios import (
     avro_all_types_scenario,
     avro_file_with_double_as_number_scenario,
@@ -248,10 +247,8 @@ def test_file_based_discover(capsys: CaptureFixture[str], tmp_path: PosixPath, s
 
 @pytest.mark.parametrize("scenario", read_scenarios, ids=[s.name for s in read_scenarios])
 @freeze_time("2023-06-09T00:00:00Z")
-def test_file_based_read(
-    capsys: CaptureFixture[str], caplog: LogCaptureFixture, tmp_path: PosixPath, scenario: TestScenario[AbstractSource]
-) -> None:
-    verify_read(capsys, caplog, tmp_path, scenario)
+def test_file_based_read(scenario: TestScenario[AbstractSource]) -> None:
+    verify_read(scenario)
 
 
 @pytest.mark.parametrize("scenario", spec_scenarios, ids=[c.name for c in spec_scenarios])

unit_tests/sources/file_based/test_scenarios.py

@@ -11,11 +11,12 @@ import pytest
 from _pytest.capture import CaptureFixture
 from _pytest.reports import ExceptionInfo
 from airbyte_cdk.entrypoint import launch
-from airbyte_cdk.logger import AirbyteLogFormatter
 from airbyte_cdk.models import AirbyteAnalyticsTraceMessage, SyncMode
 from airbyte_cdk.sources import AbstractSource
+from airbyte_cdk.test.entrypoint_wrapper import EntrypointOutput
+from airbyte_cdk.test.entrypoint_wrapper import read as entrypoint_read
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
-from pytest import LogCaptureFixture
+from airbyte_protocol.models import AirbyteLogMessage, AirbyteMessage, ConfiguredAirbyteCatalog
 from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenario
 
 
@@ -37,58 +38,51 @@ def verify_discover(capsys: CaptureFixture[str], tmp_path: PosixPath, scenario:
             _verify_expected_logs(logs, discover_logs)
 
 
-def verify_read(
-    capsys: CaptureFixture[str], caplog: LogCaptureFixture, tmp_path: PosixPath, scenario: TestScenario[AbstractSource]
-) -> None:
-    caplog.handler.setFormatter(AirbyteLogFormatter())
+def verify_read(scenario: TestScenario[AbstractSource]) -> None:
     if scenario.incremental_scenario_config:
-        run_test_read_incremental(capsys, caplog, tmp_path, scenario)
+        run_test_read_incremental(scenario)
     else:
-        run_test_read_full_refresh(capsys, caplog, tmp_path, scenario)
+        run_test_read_full_refresh(scenario)
 
 
-def run_test_read_full_refresh(
-    capsys: CaptureFixture[str], caplog: LogCaptureFixture, tmp_path: PosixPath, scenario: TestScenario[AbstractSource]
-) -> None:
+def run_test_read_full_refresh(scenario: TestScenario[AbstractSource]) -> None:
     expected_exc, expected_msg = scenario.expected_read_error
     if expected_exc:
         with pytest.raises(expected_exc) as exc:  # noqa
-            read(capsys, caplog, tmp_path, scenario)
+            read(scenario)
         if expected_msg:
             assert expected_msg in get_error_message_from_exc(exc)
     else:
-        output = read(capsys, caplog, tmp_path, scenario)
+        output = read(scenario)
         _verify_read_output(output, scenario)
 
 
-def run_test_read_incremental(
-    capsys: CaptureFixture[str], caplog: LogCaptureFixture, tmp_path: PosixPath, scenario: TestScenario[AbstractSource]
-) -> None:
+def run_test_read_incremental(scenario: TestScenario[AbstractSource]) -> None:
     expected_exc, expected_msg = scenario.expected_read_error
     if expected_exc:
         with pytest.raises(expected_exc):
-            read_with_state(capsys, caplog, tmp_path, scenario)
+            read_with_state(scenario)
     else:
-        output = read_with_state(capsys, caplog, tmp_path, scenario)
+        output = read_with_state(scenario)
         _verify_read_output(output, scenario)
 
 
-def _verify_read_output(output: Dict[str, Any], scenario: TestScenario[AbstractSource]) -> None:
-    records, logs = output["records"], output["logs"]
-    logs = [log for log in logs if log.get("level") in scenario.log_levels]
+def _verify_read_output(output: EntrypointOutput, scenario: TestScenario[AbstractSource]) -> None:
+    records, log_messages = output.records_and_state_messages, output.logs
+    logs = [message.log for message in log_messages if message.log.level.value in scenario.log_levels]
     expected_records = scenario.expected_records
     assert len(records) == len(expected_records)
     for actual, expected in zip(records, expected_records):
-        if "record" in actual:
-            assert len(actual["record"]["data"]) == len(expected["data"])
-            for key, value in actual["record"]["data"].items():
+        if actual.record:
+            assert len(actual.record.data) == len(expected["data"])
+            for key, value in actual.record.data.items():
                 if isinstance(value, float):
                     assert math.isclose(value, expected["data"][key], abs_tol=1e-04)
                 else:
                     assert value == expected["data"][key]
-            assert actual["record"]["stream"] == expected["stream"]
-        elif "state" in actual:
-            assert actual["state"]["data"] == expected
+            assert actual.record.stream == expected["stream"]
+        elif actual.state:
+            assert actual.state.data == expected
 
     if scenario.expected_logs:
         read_logs = scenario.expected_logs.get("read")
@@ -96,25 +90,25 @@ def _verify_read_output(output: Dict[str, Any], scenario: TestScenario[AbstractS
         _verify_expected_logs(logs, read_logs)
 
     if scenario.expected_analytics:
-        analytics = output["analytics"]
+        analytics = output.analytics_messages
 
         _verify_analytics(analytics, scenario.expected_analytics)
 
 
-def _verify_analytics(analytics: List[Dict[str, Any]], expected_analytics: Optional[List[AirbyteAnalyticsTraceMessage]]) -> None:
+def _verify_analytics(analytics: List[AirbyteMessage], expected_analytics: Optional[List[AirbyteAnalyticsTraceMessage]]) -> None:
     if expected_analytics:
         for actual, expected in zip(analytics, expected_analytics):
-            actual_type, actual_value = actual["type"], actual["value"]
+            actual_type, actual_value = actual.trace.analytics.type, actual.trace.analytics.value
             expected_type = expected.type
             expected_value = expected.value
             assert actual_type == expected_type
             assert actual_value == expected_value
 
 
-def _verify_expected_logs(logs: List[Dict[str, Any]], expected_logs: Optional[List[Mapping[str, Any]]]) -> None:
+def _verify_expected_logs(logs: List[AirbyteLogMessage], expected_logs: Optional[List[Mapping[str, Any]]]) -> None:
     if expected_logs:
         for actual, expected in zip(logs, expected_logs):
-            actual_level, actual_message = actual["level"], actual["message"]
+            actual_level, actual_message = actual.level.value, actual.message
             expected_level = expected["level"]
             expected_message = expected["message"]
             assert actual_level == expected_level
@@ -172,55 +166,21 @@ def discover(capsys: CaptureFixture[str], tmp_path: PosixPath, scenario: TestSce
     }
 
 
-def read(
-    capsys: CaptureFixture[str], caplog: LogCaptureFixture, tmp_path: PosixPath, scenario: TestScenario[AbstractSource]
-) -> Dict[str, Any]:
-    with caplog.handler.stream as logger_stream:
-        launch(
-            scenario.source,
-            [
-                "read",
-                "--config",
-                make_file(tmp_path / "config.json", scenario.config),
-                "--catalog",
-                make_file(tmp_path / "catalog.json", scenario.configured_catalog(SyncMode.full_refresh)),
-            ],
-        )
-        captured = capsys.readouterr().out.splitlines() + logger_stream.getvalue().split("\n")[:-1]
-
-        return {
-            "records": [msg for msg in (json.loads(line) for line in captured) if msg["type"] == "RECORD"],
-            "logs": [msg["log"] for msg in (json.loads(line) for line in captured) if msg["type"] == "LOG"],
-            "analytics": [
-                msg["trace"]["analytics"]
-                for msg in (json.loads(line) for line in captured)
-                if msg["type"] == "TRACE" and msg["trace"]["type"] == "ANALYTICS"
-            ],
-        }
-
-
-def read_with_state(
-    capsys: CaptureFixture[str], caplog: LogCaptureFixture, tmp_path: PosixPath, scenario: TestScenario[AbstractSource]
-) -> Dict[str, List[Any]]:
-    launch(
+def read(scenario: TestScenario[AbstractSource]) -> EntrypointOutput:
+    return entrypoint_read(
         scenario.source,
-        [
-            "read",
-            "--config",
-            make_file(tmp_path / "config.json", scenario.config),
-            "--catalog",
-            make_file(tmp_path / "catalog.json", scenario.configured_catalog(SyncMode.incremental)),
-            "--state",
-            make_file(tmp_path / "state.json", scenario.input_state()),
-        ],
+        scenario.config,
+        ConfiguredAirbyteCatalog.parse_obj(scenario.configured_catalog(SyncMode.full_refresh)),
+    )
+
+
+def read_with_state(scenario: TestScenario[AbstractSource]) -> EntrypointOutput:
+    return entrypoint_read(
+        scenario.source,
+        scenario.config,
+        ConfiguredAirbyteCatalog.parse_obj(scenario.configured_catalog(SyncMode.incremental)),
+        scenario.input_state(),
     )
-    captured = capsys.readouterr()
-    logs = caplog.records
-    return {
-        "records": [msg for msg in (json.loads(line) for line in captured.out.splitlines()) if msg["type"] in ("RECORD", "STATE")],
-        "logs": [msg["log"] for msg in (json.loads(line) for line in captured.out.splitlines()) if msg["type"] == "LOG"]
-        + [{"level": log.levelname, "message": log.message} for log in logs],
-    }
 
 
 def make_file(path: Path, file_contents: Optional[Union[Mapping[str, Any], List[Mapping[str, Any]]]]) -> str:

unit_tests/sources/streams/concurrent/scenarios/test_concurrent_scenarios.py

@@ -7,7 +7,6 @@ from pathlib import PosixPath
 import pytest
 from _pytest.capture import CaptureFixture
 from freezegun import freeze_time
-from pytest import LogCaptureFixture
 from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenario
 from unit_tests.sources.file_based.test_scenarios import verify_discover, verify_read
 from unit_tests.sources.streams.concurrent.scenarios.incremental_scenarios import (
@@ -68,8 +67,8 @@ scenarios = [
 
 @pytest.mark.parametrize("scenario", scenarios, ids=[s.name for s in scenarios])
 @freeze_time("2023-06-09T00:00:00Z")
-def test_concurrent_read(capsys: CaptureFixture[str], caplog: LogCaptureFixture, tmp_path: PosixPath, scenario: TestScenario) -> None:
-    verify_read(capsys, caplog, tmp_path, scenario)
+def test_concurrent_read(scenario: TestScenario) -> None:
+    verify_read(scenario)
 
 
 @pytest.mark.parametrize("scenario", scenarios, ids=[s.name for s in scenarios])