diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b22..de006e5a2c 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -415,10 +415,11 @@ def update( def _update_current_window(self) -> None: now = datetime.datetime.now() - if now > self._next_reset_ts: - logger.debug("started new window, %s calls available now", self._call_limit) - self._next_reset_ts = self._next_reset_ts + self._offset + if now >= self._next_reset_ts: + elapsed_periods = (now - self._next_reset_ts) // self._offset + 1 + self._next_reset_ts = self._next_reset_ts + self._offset * elapsed_periods self._calls_num = 0 + logger.debug("started new window, %s calls available now", self._call_limit) class MovingWindowCallRatePolicy(BaseCallRatePolicy): diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe5737..b0cbc6f815 100644 --- a/unit_tests/sources/streams/test_call_rate.py +++ b/unit_tests/sources/streams/test_call_rate.py @@ -9,6 +9,7 @@ import pytest import requests +from freezegun import freeze_time from requests import Request from airbyte_cdk.models import SyncMode @@ -249,6 +250,72 @@ def test_update_available_calls(self, mocker): with pytest.raises(CallRateLimitHit): policy.try_acquire(mocker.Mock(), weight=3) + def test_multiple_missed_periods_enforce_limit(self, mocker): + base = datetime(2024, 1, 1, 12, 0, 0) + policy = FixedWindowCallRatePolicy( + matchers=[], + next_reset_ts=base - timedelta(minutes=3), + period=timedelta(minutes=1), + call_limit=2, + ) + + with freeze_time(base): + policy.try_acquire(mocker.Mock(), weight=1) + policy.try_acquire(mocker.Mock(), weight=1) + with pytest.raises(CallRateLimitHit): + policy.try_acquire(mocker.Mock(), weight=1) + + assert policy._next_reset_ts == base + timedelta(minutes=1) + + def test_call_allowed_exactly_at_reset_boundary(self, mocker): + base = datetime(2024, 1, 1, 12, 0, 0) + policy = FixedWindowCallRatePolicy( + matchers=[], + next_reset_ts=base + timedelta(minutes=1), + period=timedelta(minutes=1), + call_limit=1, + ) + + with freeze_time(base) as frozen: + policy.try_acquire(mocker.Mock(), weight=1) + with pytest.raises(CallRateLimitHit): + policy.try_acquire(mocker.Mock(), weight=1) + + frozen.move_to(base + timedelta(minutes=1)) + policy.try_acquire(mocker.Mock(), weight=1) + + def test_alignment_preserved_after_skipping_periods(self, mocker): + base = datetime(2024, 1, 1, 12, 0, 0) + policy = FixedWindowCallRatePolicy( + matchers=[], + next_reset_ts=base, + period=timedelta(hours=1), + call_limit=5, + ) + + with freeze_time(base + timedelta(hours=2, minutes=30)): + policy.try_acquire(mocker.Mock(), weight=1) + + assert policy._next_reset_ts == base + timedelta(hours=3) + + def test_single_period_rollover(self, mocker): + base = datetime(2024, 1, 1, 12, 0, 0) + policy = FixedWindowCallRatePolicy( + matchers=[], + next_reset_ts=base + timedelta(minutes=1), + period=timedelta(minutes=1), + call_limit=1, + ) + + with freeze_time(base) as frozen: + policy.try_acquire(mocker.Mock(), weight=1) + + frozen.move_to(base + timedelta(minutes=1, seconds=30)) + policy.try_acquire(mocker.Mock(), weight=1) + assert policy._next_reset_ts == base + timedelta(minutes=2) + with pytest.raises(CallRateLimitHit): + policy.try_acquire(mocker.Mock(), weight=1) + class TestMovingWindowCallRatePolicy: def test_no_rates(self):