airbyte-cdk 0.72.1__py3-none-any.whl → 6.17.1.dev0__py3-none-any.whl

unit_tests/test/test_entrypoint_wrapper.py

@@ -1,240 +0,0 @@
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-
-import json
-import logging
-import os
-from typing import Any, Iterator, List, Mapping
-from unittest import TestCase
-from unittest.mock import Mock, patch
-
-from airbyte_cdk.sources.abstract_source import AbstractSource
-from airbyte_cdk.test.entrypoint_wrapper import read
-from airbyte_cdk.test.state_builder import StateBuilder
-from airbyte_protocol.models import (
-    AirbyteAnalyticsTraceMessage,
-    AirbyteErrorTraceMessage,
-    AirbyteLogMessage,
-    AirbyteMessage,
-    AirbyteRecordMessage,
-    AirbyteStateBlob,
-    AirbyteStateMessage,
-    AirbyteStreamState,
-    AirbyteStreamStatus,
-    AirbyteStreamStatusTraceMessage,
-    AirbyteTraceMessage,
-    ConfiguredAirbyteCatalog,
-    Level,
-    StreamDescriptor,
-    TraceType,
-    Type,
-)
-
-
-def _a_state_message(stream_name: str, stream_state: Mapping[str, Any]) -> AirbyteMessage:
-    return AirbyteMessage(type=Type.STATE, state=AirbyteStateMessage(stream=AirbyteStreamState(stream_descriptor=StreamDescriptor(name=stream_name), stream_state=AirbyteStateBlob(**stream_state))))
-
-
-def _a_status_message(stream_name: str, status: AirbyteStreamStatus) -> AirbyteMessage:
-    return AirbyteMessage(
-        type=Type.TRACE,
-        trace=AirbyteTraceMessage(
-            type=TraceType.STREAM_STATUS,
-            emitted_at=0,
-            stream_status=AirbyteStreamStatusTraceMessage(
-                stream_descriptor=StreamDescriptor(name=stream_name),
-                status=status,
-            ),
-        ),
-    )
-
-
-_A_RECORD = AirbyteMessage(
-    type=Type.RECORD, record=AirbyteRecordMessage(stream="stream", data={"record key": "record value"}, emitted_at=0)
-)
-_A_STATE_MESSAGE = _a_state_message("stream_name", {"state key": "state value for _A_STATE_MESSAGE"})
-_A_LOG = AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is an Airbyte log message"))
-_AN_ERROR_MESSAGE = AirbyteMessage(
-    type=Type.TRACE,
-    trace=AirbyteTraceMessage(
-        type=TraceType.ERROR,
-        emitted_at=0,
-        error=AirbyteErrorTraceMessage(message="AirbyteErrorTraceMessage message"),
-    ),
-)
-_AN_ANALYTIC_MESSAGE = AirbyteMessage(
-    type=Type.TRACE,
-    trace=AirbyteTraceMessage(
-        type=TraceType.ANALYTICS,
-        emitted_at=0,
-        analytics=AirbyteAnalyticsTraceMessage(type="an analytic type", value="an analytic value"),
-    ),
-)
-
-_A_STREAM_NAME = "a stream name"
-_A_CONFIG = {"config_key": "config_value"}
-_A_CATALOG = ConfiguredAirbyteCatalog.parse_obj(
-    {
-        "streams": [
-            {
-                "stream": {
-                    "name": "a_stream_name",
-                    "json_schema": {},
-                    "supported_sync_modes": ["full_refresh"],
-                },
-                "sync_mode": "full_refresh",
-                "destination_sync_mode": "append",
-            }
-        ]
-    }
-)
-_A_STATE = StateBuilder().with_stream_state(_A_STREAM_NAME, {"state_key": "state_value"}).build()
-_A_LOG_MESSAGE = "a log message"
-
-
-def _to_entrypoint_output(messages: List[AirbyteMessage]) -> Iterator[str]:
-    return (message.json(exclude_unset=True) for message in messages)
-
-
-def _a_mocked_source() -> AbstractSource:
-    source = Mock(spec=AbstractSource)
-    source.message_repository = None
-    return source
-
-
-def _validate_tmp_json_file(expected, file_path) -> None:
-    with open(file_path) as file:
-        assert json.load(file) == expected
-
-
-def _validate_tmp_catalog(expected, file_path) -> None:
-    assert ConfiguredAirbyteCatalog.parse_file(file_path) == expected
-
-
-def _create_tmp_file_validation(entrypoint, expected_config, expected_catalog, expected_state):
-    def _validate_tmp_files(self):
-        _validate_tmp_json_file(expected_config, entrypoint.return_value.parse_args.call_args.args[0][2])
-        _validate_tmp_catalog(expected_catalog, entrypoint.return_value.parse_args.call_args.args[0][4])
-        _validate_tmp_json_file(expected_state, entrypoint.return_value.parse_args.call_args.args[0][6])
-        return entrypoint.return_value.run.return_value
-
-    return _validate_tmp_files
-
-
-class EntrypointWrapperTest(TestCase):
-    def setUp(self) -> None:
-        self._a_source = _a_mocked_source()
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_when_read_then_ensure_parameters(self, entrypoint):
-        entrypoint.return_value.run.side_effect = _create_tmp_file_validation(entrypoint, _A_CONFIG, _A_CATALOG, _A_STATE)
-
-        read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-
-        entrypoint.assert_called_once_with(self._a_source)
-        entrypoint.return_value.run.assert_called_once_with(entrypoint.return_value.parse_args.return_value)
-        assert entrypoint.return_value.parse_args.call_count == 1
-        assert entrypoint.return_value.parse_args.call_args.args[0][0] == "read"
-        assert entrypoint.return_value.parse_args.call_args.args[0][1] == "--config"
-        assert entrypoint.return_value.parse_args.call_args.args[0][3] == "--catalog"
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_when_read_then_ensure_files_are_temporary(self, entrypoint):
-        read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-
-        assert not os.path.exists(entrypoint.return_value.parse_args.call_args.args[0][2])
-        assert not os.path.exists(entrypoint.return_value.parse_args.call_args.args[0][4])
-        assert not os.path.exists(entrypoint.return_value.parse_args.call_args.args[0][6])
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_logging_during_run_when_read_then_output_has_logs(self, entrypoint):
-        def _do_some_logging(self):
-            logging.getLogger("any logger").info(_A_LOG_MESSAGE)
-            return entrypoint.return_value.run.return_value
-
-        entrypoint.return_value.run.side_effect = _do_some_logging
-
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-
-        assert len(output.logs) == 1
-        assert output.logs[0].log.message == _A_LOG_MESSAGE
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_record_when_read_then_output_has_record(self, entrypoint):
-        entrypoint.return_value.run.return_value = _to_entrypoint_output([_A_RECORD])
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert output.records == [_A_RECORD]
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_state_message_when_read_then_output_has_state_message(self, entrypoint):
-        entrypoint.return_value.run.return_value = _to_entrypoint_output([_A_STATE_MESSAGE])
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert output.state_messages == [_A_STATE_MESSAGE]
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_state_message_and_records_when_read_then_output_has_records_and_state_message(self, entrypoint):
-        entrypoint.return_value.run.return_value = _to_entrypoint_output([_A_RECORD, _A_STATE_MESSAGE])
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert output.records_and_state_messages == [_A_RECORD, _A_STATE_MESSAGE]
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_many_state_messages_and_records_when_read_then_output_has_records_and_state_message(self, entrypoint):
-        state_value = {"state_key": "last state value"}
-        last_emitted_state = AirbyteStreamState(stream_descriptor=StreamDescriptor(name="stream_name"), stream_state=AirbyteStateBlob(**state_value))
-        entrypoint.return_value.run.return_value = _to_entrypoint_output([_A_STATE_MESSAGE, _a_state_message("stream_name", state_value)])
-
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-
-        assert output.most_recent_state == last_emitted_state
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_log_when_read_then_output_has_log(self, entrypoint):
-        entrypoint.return_value.run.return_value = _to_entrypoint_output([_A_LOG])
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert output.logs == [_A_LOG]
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_trace_message_when_read_then_output_has_trace_messages(self, entrypoint):
-        entrypoint.return_value.run.return_value = _to_entrypoint_output([_AN_ANALYTIC_MESSAGE])
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert output.analytics_messages == [_AN_ANALYTIC_MESSAGE]
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_stream_statuses_when_read_then_return_statuses(self, entrypoint):
-        status_messages = [
-            _a_status_message(_A_STREAM_NAME, AirbyteStreamStatus.STARTED),
-            _a_status_message(_A_STREAM_NAME, AirbyteStreamStatus.COMPLETE),
-        ]
-        entrypoint.return_value.run.return_value = _to_entrypoint_output(status_messages)
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert output.get_stream_statuses(_A_STREAM_NAME) == [AirbyteStreamStatus.STARTED, AirbyteStreamStatus.COMPLETE]
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_stream_statuses_for_many_streams_when_read_then_filter_other_streams(self, entrypoint):
-        status_messages = [
-            _a_status_message(_A_STREAM_NAME, AirbyteStreamStatus.STARTED),
-            _a_status_message("another stream name", AirbyteStreamStatus.INCOMPLETE),
-            _a_status_message(_A_STREAM_NAME, AirbyteStreamStatus.COMPLETE),
-        ]
-        entrypoint.return_value.run.return_value = _to_entrypoint_output(status_messages)
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert len(output.get_stream_statuses(_A_STREAM_NAME)) == 2
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.print", create=True)
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_unexpected_exception_when_read_then_print(self, entrypoint, print_mock):
-        entrypoint.return_value.run.side_effect = ValueError("This error should be printed")
-        read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert print_mock.call_count > 0
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.print", create=True)
-    @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 be printed")
-        read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE, expecting_exception=True)
-        assert print_mock.call_count == 0
-
-    @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
-    def test_given_uncaught_exception_when_read_then_output_has_error(self, entrypoint):
-        entrypoint.return_value.run.side_effect = ValueError("This error should be printed")
-        output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
-        assert output.errors

unit_tests/utils/test_datetime_format_inferrer.py

@@ -1,60 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-from typing import Dict, List
-
-import pytest
-from airbyte_cdk.models.airbyte_protocol import AirbyteRecordMessage
-from airbyte_cdk.utils.datetime_format_inferrer import DatetimeFormatInferrer
-
-NOW = 1234567
-
-
-@pytest.mark.parametrize(
-    "test_name,input_records,expected_candidate_fields",
-    [
-        ("empty", [], {}),
-        ("simple_match", [{"d": "2022-02-03"}], {"d": "%Y-%m-%d"}),
-        ("timestamp_match_integer", [{"d": 1686058051}], {"d": "%s"}),
-        ("timestamp_match_string", [{"d": "1686058051"}], {"d": "%s"}),
-        ("timestamp_ms_match_integer", [{"d": 1686058051000}], {"d": "%ms"}),
-        ("timestamp_ms_match_string", [{"d": "1686058051000"}], {"d": "%ms"}),
-        ("timestamp_no_match_integer", [{"d": 99}], {}),
-        ("timestamp_no_match_string", [{"d": "99999999999999999999"}], {}),
-        ("timestamp_overflow", [{"d": f"{10**100}_100"}], {}),  # this case was previously causing OverflowError hence this test
-        ("simple_no_match", [{"d": "20220203"}], {}),
-        ("multiple_match", [{"d": "2022-02-03", "e": "2022-02-03"}], {"d": "%Y-%m-%d", "e": "%Y-%m-%d"}),
-        (
-            "multiple_no_match",
-            [{"d": "20220203", "r": "ccc", "e": {"something-else": "2023-03-03"}, "s": ["2023-03-03"], "x": False, "y": 123}],
-            {},
-        ),
-        ("format_1", [{"d": "2022-02-03"}], {"d": "%Y-%m-%d"}),
-        ("format_2", [{"d": "2022-02-03 12:34:56"}], {"d": "%Y-%m-%d %H:%M:%S"}),
-        ("format_3", [{"d": "2022-02-03T12:34:56Z"}], {"d": "%Y-%m-%dT%H:%M:%SZ"}),
-        ("format_4 1", [{"d": "2022-02-03T12:34:56.000Z"}], {"d": "%Y-%m-%dT%H:%M:%S.%fZ"}),
-        ("format_4 2", [{"d": "2022-02-03T12:34:56.000000Z"}], {"d": "%Y-%m-%dT%H:%M:%S.%fZ"}),
-        ("format_5", [{"d": "2022-02-03 12:34:56.123456+00:00"}], {"d": "%Y-%m-%d %H:%M:%S.%f%z"}),
-        ("format_5 2", [{"d": "2022-02-03 12:34:56.123456+02:00"}], {"d": "%Y-%m-%d %H:%M:%S.%f%z"}),
-        ("format_6", [{"d": "2022-02-03T12:34:56.123456+0000"}], {"d": "%Y-%m-%dT%H:%M:%S.%f%z"}),
-        ("format_6 2", [{"d": "2022-02-03T12:34:56.123456+00:00"}], {"d": "%Y-%m-%dT%H:%M:%S.%f%z"}),
-        ("format_6 3", [{"d": "2022-02-03T12:34:56.123456-03:00"}], {"d": "%Y-%m-%dT%H:%M:%S.%f%z"}),
-        ("format_7", [{"d": "03/02/2022 12:34"}], {"d": "%d/%m/%Y %H:%M"}),
-        ("format_8", [{"d": "2022-02"}], {"d": "%Y-%m"}),
-        ("format_9", [{"d": "03-02-2022"}], {"d": "%d-%m-%Y"}),
-        ("limit_down", [{"d": "2022-02-03", "x": "2022-02-03"}, {"d": "2022-02-03", "x": "another thing"}], {"d": "%Y-%m-%d"}),
-        ("limit_down all", [{"d": "2022-02-03", "x": "2022-02-03"}, {"d": "also another thing", "x": "another thing"}], {}),
-        ("limit_down empty", [{"d": "2022-02-03", "x": "2022-02-03"}, {}], {}),
-        ("limit_down unsupported type", [{"d": "2022-02-03"}, {"d": False}], {}),
-        ("limit_down complex type", [{"d": "2022-02-03"}, {"d": {"date": "2022-03-03"}}], {}),
-        ("limit_down different format", [{"d": "2022-02-03"}, {"d": 1686058051}], {}),
-        ("limit_down different format", [{"d": "2022-02-03"}, {"d": "2022-02-03T12:34:56.000000Z"}], {}),
-        ("no scope expand", [{}, {"d": "2022-02-03"}], {}),
-    ],
-)
-def test_schema_inferrer(test_name, input_records: List, expected_candidate_fields: Dict[str, str]):
-    inferrer = DatetimeFormatInferrer()
-    for record in input_records:
-        inferrer.accumulate(AirbyteRecordMessage(stream="abc", data=record, emitted_at=NOW))
-    assert inferrer.get_inferred_datetime_formats() == expected_candidate_fields

unit_tests/utils/test_mapping_helpers.py

@@ -1,54 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-import pytest
-from airbyte_cdk.utils.mapping_helpers import combine_mappings
-
-
-def test_basic_merge():
-    mappings = [{"a": 1}, {"b": 2}, {"c": 3}, {}]
-    result = combine_mappings(mappings)
-    assert result == {"a": 1, "b": 2, "c": 3}
-
-
-def test_combine_with_string():
-    mappings = [{"a": 1}, "option"]
-    with pytest.raises(ValueError, match="Cannot combine multiple options if one is a string"):
-        combine_mappings(mappings)
-
-
-def test_overlapping_keys():
-    mappings = [{"a": 1, "b": 2}, {"b": 3}]
-    with pytest.raises(ValueError, match="Duplicate keys found"):
-        combine_mappings(mappings)
-
-
-def test_multiple_strings():
-    mappings = ["option1", "option2"]
-    with pytest.raises(ValueError, match="Cannot combine multiple string options"):
-        combine_mappings(mappings)
-
-
-def test_handle_none_values():
-    mappings = [{"a": 1}, None, {"b": 2}]
-    result = combine_mappings(mappings)
-    assert result == {"a": 1, "b": 2}
-
-
-def test_empty_mappings():
-    mappings = []
-    result = combine_mappings(mappings)
-    assert result == {}
-
-
-def test_single_mapping():
-    mappings = [{"a": 1}]
-    result = combine_mappings(mappings)
-    assert result == {"a": 1}
-
-
-def test_combine_with_string_and_empty_mappings():
-    mappings = ["option", {}]
-    result = combine_mappings(mappings)
-    assert result == "option"

unit_tests/utils/test_message_utils.py

@@ -1,91 +0,0 @@
-# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
-
-import pytest
-from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
-from airbyte_cdk.utils.message_utils import get_stream_descriptor
-from airbyte_protocol.models import (
-    AirbyteControlConnectorConfigMessage,
-    AirbyteControlMessage,
-    AirbyteMessage,
-    AirbyteRecordMessage,
-    AirbyteStateBlob,
-    AirbyteStateMessage,
-    AirbyteStateStats,
-    AirbyteStateType,
-    AirbyteStreamState,
-    OrchestratorType,
-    StreamDescriptor,
-    Type,
-)
-
-
-def test_get_record_message_stream_descriptor():
-    message = AirbyteMessage(
-        type=Type.RECORD,
-        record=AirbyteRecordMessage(
-            stream="test_stream",
-            namespace="test_namespace",
-            data={"id": "12345"},
-            emitted_at=1,
-        ),
-    )
-    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace="test_namespace")
-    assert get_stream_descriptor(message) == expected_descriptor
-
-
-def test_get_record_message_stream_descriptor_no_namespace():
-    message = AirbyteMessage(
-        type=Type.RECORD,
-        record=AirbyteRecordMessage(
-            stream="test_stream", data={"id": "12345"}, emitted_at=1
-        ),
-    )
-    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace=None)
-    assert get_stream_descriptor(message) == expected_descriptor
-
-
-def test_get_state_message_stream_descriptor():
-    message = AirbyteMessage(
-        type=Type.STATE,
-        state=AirbyteStateMessage(
-            type=AirbyteStateType.STREAM,
-            stream=AirbyteStreamState(
-                stream_descriptor=StreamDescriptor(
-                    name="test_stream", namespace="test_namespace"
-                ),
-                stream_state=AirbyteStateBlob(updated_at="2024-02-02"),
-            ),
-            sourceStats=AirbyteStateStats(recordCount=27.0),
-        ),
-    )
-    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace="test_namespace")
-    assert get_stream_descriptor(message) == expected_descriptor
-
-
-def test_get_state_message_stream_descriptor_no_namespace():
-    message = AirbyteMessage(
-        type=Type.STATE,
-        state=AirbyteStateMessage(
-            type=AirbyteStateType.STREAM,
-            stream=AirbyteStreamState(
-                stream_descriptor=StreamDescriptor(name="test_stream"),
-                stream_state=AirbyteStateBlob(updated_at="2024-02-02"),
-            ),
-            sourceStats=AirbyteStateStats(recordCount=27.0),
-        ),
-    )
-    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace=None)
-    assert get_stream_descriptor(message) == expected_descriptor
-
-
-def test_get_other_message_stream_descriptor_fails():
-    message = AirbyteMessage(
-        type=Type.CONTROL,
-        control=AirbyteControlMessage(
-            type=OrchestratorType.CONNECTOR_CONFIG,
-            emitted_at=10,
-            connectorConfig=AirbyteControlConnectorConfigMessage(config={"any config": "a config value"}),
-        ),
-    )
-    with pytest.raises(NotImplementedError):
-        get_stream_descriptor(message)

unit_tests/utils/test_rate_limiting.py

@@ -1,26 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-import pytest
-from airbyte_cdk.sources.streams.http.rate_limiting import default_backoff_handler
-from requests import exceptions
-
-
-def helper_with_exceptions(exception_type):
-    raise exception_type
-
-
-@pytest.mark.parametrize(
-    "max_tries, max_time, factor, exception_to_raise",
-    [
-        (1, None, 1, exceptions.ConnectTimeout),
-        (1, 1, 0, exceptions.ReadTimeout),
-        (2, 2, 1, exceptions.ConnectionError),
-        (3, 3, 1, exceptions.ChunkedEncodingError),
-    ],
-)
-def test_default_backoff_handler(max_tries: int, max_time: int, factor: int, exception_to_raise: Exception):
-    backoff_handler = default_backoff_handler(max_tries=max_tries, max_time=max_time, factor=factor)(helper_with_exceptions)
-    with pytest.raises(exception_to_raise):
-        backoff_handler(exception_to_raise)