airbyte-cdk 0.72.1__py3-none-any.whl → 6.13.1.dev4107__py3-none-any.whl

unit_tests/sources/streams/test_availability_strategy.py

@@ -1,70 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-import logging
-from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union
-
-from airbyte_cdk.models import SyncMode
-from airbyte_cdk.sources import Source
-from airbyte_cdk.sources.streams import Stream
-from airbyte_cdk.sources.streams.availability_strategy import AvailabilityStrategy
-from airbyte_cdk.sources.streams.core import StreamData
-
-logger = logging.getLogger("airbyte")
-
-
-class MockStream(Stream):
-    def __init__(self, name: str) -> Stream:
-        self._name = name
-
-    @property
-    def name(self) -> str:
-        return self._name
-
-    @property
-    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
-        pass
-
-    def read_records(
-        self,
-        sync_mode: SyncMode,
-        cursor_field: List[str] = None,
-        stream_slice: Mapping[str, Any] = None,
-        stream_state: Mapping[str, Any] = None,
-    ) -> Iterable[StreamData]:
-        pass
-
-
-def test_no_availability_strategy():
-    stream_1 = MockStream("stream")
-    assert stream_1.availability_strategy is None
-
-    stream_1_is_available, _ = stream_1.check_availability(logger)
-    assert stream_1_is_available
-
-
-def test_availability_strategy():
-    class MockAvailabilityStrategy(AvailabilityStrategy):
-        def check_availability(self, stream: Stream, logger: logging.Logger, source: Optional[Source]) -> Tuple[bool, any]:
-            if stream.name == "available_stream":
-                return True, None
-            return False, f"Could not reach stream '{stream.name}'."
-
-    class MockStreamWithAvailabilityStrategy(MockStream):
-        @property
-        def availability_strategy(self) -> Optional["AvailabilityStrategy"]:
-            return MockAvailabilityStrategy()
-
-    stream_1 = MockStreamWithAvailabilityStrategy("available_stream")
-    stream_2 = MockStreamWithAvailabilityStrategy("unavailable_stream")
-
-    for stream in [stream_1, stream_2]:
-        assert isinstance(stream.availability_strategy, MockAvailabilityStrategy)
-
-    stream_1_is_available, _ = stream_1.check_availability(logger)
-    assert stream_1_is_available
-
-    stream_2_is_available, reason = stream_2.check_availability(logger)
-    assert not stream_2_is_available
-    assert "Could not reach stream 'unavailable_stream'" in reason

unit_tests/sources/streams/test_call_rate.py

@@ -1,300 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-import os
-import tempfile
-import time
-from datetime import datetime, timedelta
-from typing import Iterable, Mapping
-
-import pytest
-from airbyte_cdk.models import SyncMode
-from airbyte_cdk.sources.streams.call_rate import (
-    APIBudget,
-    CallRateLimitHit,
-    FixedWindowCallRatePolicy,
-    HttpRequestMatcher,
-    MovingWindowCallRatePolicy,
-    Rate,
-    UnlimitedCallRatePolicy,
-)
-from airbyte_cdk.sources.streams.http import HttpStream
-from airbyte_cdk.sources.streams.http.requests_native_auth import TokenAuthenticator
-from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH
-from requests import Request
-
-
-class StubDummyHttpStream(HttpStream):
-    url_base = "https://test_base_url.com"
-    primary_key = "some_key"
-
-    def next_page_token(self, *args, **kwargs):
-        return True  # endless pages
-
-    def path(self, **kwargs) -> str:
-        return ""
-
-    def parse_response(self, *args, **kwargs) -> Iterable[Mapping]:
-        yield {"data": "some_data"}
-
-
-class StubDummyCacheHttpStream(StubDummyHttpStream):
-    use_cache = True
-
-
-@pytest.fixture(name="enable_cache")
-def enable_cache_fixture():
-    prev_cache_path = os.environ.get(ENV_REQUEST_CACHE_PATH)
-    with tempfile.TemporaryDirectory() as temp_dir:
-        os.environ[ENV_REQUEST_CACHE_PATH] = temp_dir
-        yield
-
-    if prev_cache_path is not None:
-        os.environ[ENV_REQUEST_CACHE_PATH] = prev_cache_path
-
-
-class TestHttpRequestMatcher:
-    try_all_types_of_requests = pytest.mark.parametrize(
-        "request_factory",
-        [Request, lambda *args, **kwargs: Request(*args, **kwargs).prepare()],
-    )
-
-    @try_all_types_of_requests
-    def test_url(self, request_factory):
-        matcher = HttpRequestMatcher(url="http://some_url/")
-        assert not matcher(request_factory(url="http://some_wrong_url"))
-        assert matcher(request_factory(url="http://some_url"))
-
-    @try_all_types_of_requests
-    def test_method(self, request_factory):
-        matcher = HttpRequestMatcher(method="GET")
-        assert not matcher(request_factory(url="http://some_url"))
-        assert not matcher(request_factory(url="http://some_url", method="POST"))
-        assert matcher(request_factory(url="http://some_url", method="GET"))
-
-    @try_all_types_of_requests
-    def test_params(self, request_factory):
-        matcher = HttpRequestMatcher(params={"param1": 10, "param2": 15})
-        assert not matcher(request_factory(url="http://some_url/"))
-        assert not matcher(request_factory(url="http://some_url/", params={"param1": 10, "param3": 100}))
-        assert not matcher(request_factory(url="http://some_url/", params={"param1": 10, "param2": 10}))
-        assert matcher(request_factory(url="http://some_url/", params={"param1": 10, "param2": 15, "param3": 100}))
-
-    @try_all_types_of_requests
-    def test_header(self, request_factory):
-        matcher = HttpRequestMatcher(headers={"header1": 10, "header2": 15})
-        assert not matcher(request_factory(url="http://some_url"))
-        assert not matcher(request_factory(url="http://some_url", headers={"header1": "10", "header3": "100"}))
-        assert not matcher(request_factory(url="http://some_url", headers={"header1": "10", "header2": "10"}))
-        assert matcher(request_factory(url="http://some_url", headers={"header1": "10", "header2": "15", "header3": "100"}))
-
-    @try_all_types_of_requests
-    def test_combination(self, request_factory):
-        matcher = HttpRequestMatcher(method="GET", url="http://some_url/", headers={"header1": 10}, params={"param2": "test"})
-        assert matcher(request_factory(method="GET", url="http://some_url", headers={"header1": "10"}, params={"param2": "test"}))
-        assert not matcher(request_factory(method="GET", url="http://some_url", headers={"header1": "10"}))
-        assert not matcher(request_factory(method="GET", url="http://some_url"))
-        assert not matcher(request_factory(url="http://some_url"))
-
-
-def test_http_request_matching(mocker):
-    """Test policy lookup based on matchers."""
-    users_policy = mocker.Mock(spec=MovingWindowCallRatePolicy)
-    groups_policy = mocker.Mock(spec=MovingWindowCallRatePolicy)
-    root_policy = mocker.Mock(spec=MovingWindowCallRatePolicy)
-
-    users_policy.matches.side_effect = HttpRequestMatcher(url="http://domain/api/users", method="GET")
-    groups_policy.matches.side_effect = HttpRequestMatcher(url="http://domain/api/groups", method="POST")
-    root_policy.matches.side_effect = HttpRequestMatcher(method="GET")
-    api_budget = APIBudget(
-        policies=[
-            users_policy,
-            groups_policy,
-            root_policy,
-        ]
-    )
-
-    api_budget.acquire_call(Request("POST", url="http://domain/unmatched_endpoint"), block=False), "unrestricted call"
-    users_policy.try_acquire.assert_not_called()
-    groups_policy.try_acquire.assert_not_called()
-    root_policy.try_acquire.assert_not_called()
-
-    users_request = Request("GET", url="http://domain/api/users")
-    api_budget.acquire_call(users_request, block=False), "first call, first matcher"
-    users_policy.try_acquire.assert_called_once_with(users_request, weight=1)
-    groups_policy.try_acquire.assert_not_called()
-    root_policy.try_acquire.assert_not_called()
-
-    api_budget.acquire_call(Request("GET", url="http://domain/api/users"), block=False), "second call, first matcher"
-    assert users_policy.try_acquire.call_count == 2
-    groups_policy.try_acquire.assert_not_called()
-    root_policy.try_acquire.assert_not_called()
-
-    group_request = Request("POST", url="http://domain/api/groups")
-    api_budget.acquire_call(group_request, block=False), "first call, second matcher"
-    assert users_policy.try_acquire.call_count == 2
-    groups_policy.try_acquire.assert_called_once_with(group_request, weight=1)
-    root_policy.try_acquire.assert_not_called()
-
-    api_budget.acquire_call(Request("POST", url="http://domain/api/groups"), block=False), "second call, second matcher"
-    assert users_policy.try_acquire.call_count == 2
-    assert groups_policy.try_acquire.call_count == 2
-    root_policy.try_acquire.assert_not_called()
-
-    any_get_request = Request("GET", url="http://domain/api/")
-    api_budget.acquire_call(any_get_request, block=False), "first call, third matcher"
-    assert users_policy.try_acquire.call_count == 2
-    assert groups_policy.try_acquire.call_count == 2
-    root_policy.try_acquire.assert_called_once_with(any_get_request, weight=1)
-
-
-class TestUnlimitedCallRatePolicy:
-    def test_try_acquire(self, mocker):
-        policy = UnlimitedCallRatePolicy(matchers=[])
-        assert policy.matches(mocker.Mock()), "should match anything"
-        policy.try_acquire(mocker.Mock(), weight=1)
-        policy.try_acquire(mocker.Mock(), weight=10)
-
-    def test_update(self):
-        policy = UnlimitedCallRatePolicy(matchers=[])
-        policy.update(available_calls=10, call_reset_ts=datetime.now())
-        policy.update(available_calls=None, call_reset_ts=datetime.now())
-        policy.update(available_calls=10, call_reset_ts=None)
-
-
-class TestFixedWindowCallRatePolicy:
-    def test_limit_rate(self, mocker):
-        policy = FixedWindowCallRatePolicy(matchers=[], next_reset_ts=datetime.now(), period=timedelta(hours=1), call_limit=100)
-        policy.try_acquire(mocker.Mock(), weight=1)
-        policy.try_acquire(mocker.Mock(), weight=20)
-        with pytest.raises(ValueError, match="Weight can not exceed the call limit"):
-            policy.try_acquire(mocker.Mock(), weight=101)
-
-        with pytest.raises(CallRateLimitHit) as exc:
-            policy.try_acquire(mocker.Mock(), weight=100 - 20 - 1 + 1)
-
-        assert exc.value.time_to_wait
-        assert exc.value.weight == 100 - 20 - 1 + 1
-        assert exc.value.item
-
-    def test_update_available_calls(self, mocker):
-        policy = FixedWindowCallRatePolicy(matchers=[], next_reset_ts=datetime.now(), period=timedelta(hours=1), call_limit=100)
-        # update to decrease number of calls available
-        policy.update(available_calls=2, call_reset_ts=None)
-        # hit the limit with weight=3
-        with pytest.raises(CallRateLimitHit):
-            policy.try_acquire(mocker.Mock(), weight=3)
-        # ok with less weight=1
-        policy.try_acquire(mocker.Mock(), weight=1)
-
-        # update to increase number of calls available, ignored
-        policy.update(available_calls=20, call_reset_ts=None)
-        # so we still hit the limit with weight=3
-        with pytest.raises(CallRateLimitHit):
-            policy.try_acquire(mocker.Mock(), weight=3)
-
-
-class TestMovingWindowCallRatePolicy:
-    def test_no_rates(self):
-        """should raise a ValueError when no rates provided"""
-        with pytest.raises(ValueError, match="The list of rates can not be empty"):
-            MovingWindowCallRatePolicy(rates=[], matchers=[])
-
-    def test_limit_rate(self):
-        """try_acquire must respect configured call rate and throw CallRateLimitHit when hit the limit."""
-        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
-
-        for i in range(10):
-            policy.try_acquire("call", weight=1), f"{i + 1} call"
-
-        with pytest.raises(CallRateLimitHit) as excinfo1:
-            policy.try_acquire("call", weight=1), "call over limit"
-        assert excinfo1.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1)
-
-        time.sleep(0.5)
-
-        with pytest.raises(CallRateLimitHit) as excinfo2:
-            policy.try_acquire("call", weight=1), "call over limit"
-        assert excinfo2.value.time_to_wait < excinfo1.value.time_to_wait, "time to wait must decrease over time"
-
-    def test_limit_rate_support_custom_weight(self):
-        """try_acquire must take into account provided weight and throw CallRateLimitHit when hit the limit."""
-        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
-
-        policy.try_acquire("call", weight=2), "1st call with weight of 2"
-        with pytest.raises(CallRateLimitHit) as excinfo:
-            policy.try_acquire("call", weight=9), "2nd call, over limit since 2 + 9 = 11 > 10"
-        assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1), "should wait 1 minute before next call"
-
-    def test_multiple_limit_rates(self):
-        """try_acquire must take into all call rates and apply stricter."""
-        policy = MovingWindowCallRatePolicy(
-            matchers=[],
-            rates=[
-                Rate(10, timedelta(minutes=10)),
-                Rate(3, timedelta(seconds=10)),
-                Rate(2, timedelta(hours=1)),
-            ],
-        )
-
-        policy.try_acquire("call", weight=2), "1 call"
-
-        with pytest.raises(CallRateLimitHit) as excinfo:
-            policy.try_acquire("call", weight=1), "1 call"
-
-        assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(3600, 0.1)
-        assert str(excinfo.value) == "Bucket for item=call with Rate limit=2/1.0h is already full"
-
-
-class TestHttpStreamIntegration:
-    def test_without_cache(self, mocker, requests_mock):
-        """Test that HttpStream will use call budget when provided"""
-        requests_mock.get(f"{StubDummyHttpStream.url_base}/", json={"data": "test"})
-
-        mocker.patch.object(MovingWindowCallRatePolicy, "try_acquire")
-
-        api_budget = APIBudget(
-            policies=[
-                MovingWindowCallRatePolicy(
-                    matchers=[HttpRequestMatcher(url=f"{StubDummyHttpStream.url_base}/", method="GET")],
-                    rates=[
-                        Rate(2, timedelta(minutes=1)),
-                    ],
-                ),
-            ]
-        )
-
-        stream = StubDummyHttpStream(api_budget=api_budget, authenticator=TokenAuthenticator(token="ABCD"))
-        records = stream.read_records(SyncMode.full_refresh)
-        for i in range(10):
-            assert next(records) == {"data": "some_data"}
-
-        assert MovingWindowCallRatePolicy.try_acquire.call_count == 10
-
-    @pytest.mark.usefixtures("enable_cache")
-    def test_with_cache(self, mocker, requests_mock):
-        """Test that HttpStream will use call budget when provided and not cached"""
-        requests_mock.get(f"{StubDummyHttpStream.url_base}/", json={"data": "test"})
-
-        mocker.patch.object(MovingWindowCallRatePolicy, "try_acquire")
-
-        api_budget = APIBudget(
-            policies=[
-                MovingWindowCallRatePolicy(
-                    matchers=[
-                        HttpRequestMatcher(url=f"{StubDummyHttpStream.url_base}/", method="GET"),
-                    ],
-                    rates=[
-                        Rate(2, timedelta(minutes=1)),
-                    ],
-                )
-            ]
-        )
-
-        stream = StubDummyCacheHttpStream(api_budget=api_budget)
-        records = stream.read_records(SyncMode.full_refresh)
-
-        for i in range(10):
-            assert next(records) == {"data": "some_data"}
-
-        assert MovingWindowCallRatePolicy.try_acquire.call_count == 1