fix: make APIBudget honor maximum_attempts_to_acquire exactly

airbyte_cdk/sources/streams/call_rate.py

@@ -555,8 +555,13 @@ def __init__(
         :param policies: list of policies in this budget
         :param maximum_attempts_to_acquire: number of attempts before throwing hit ratelimit exception, we put some big number here
          to avoid situations when many threads compete with each other for a few lots over a significant amount of time
+        :raises: ValueError if maximum_attempts_to_acquire is less than 1
         """
 
+        if maximum_attempts_to_acquire < 1:
+            raise ValueError(
+                f"maximum_attempts_to_acquire must be at least 1, got {maximum_attempts_to_acquire}"
+            )
         self._policies = policies
         self._maximum_attempts_to_acquire = maximum_attempts_to_acquire
 
@@ -628,13 +633,15 @@ def _do_acquire(
         last_exception = None
         endpoint = self._extract_endpoint(request)
         # sometimes we spend all budget before a second attempt, so we have a few more attempts
-        for attempt in range(1, self._maximum_attempts_to_acquire):
+        for attempt in range(1, self._maximum_attempts_to_acquire + 1):
             try:
                 weight = policy.get_weight(request) if isinstance(policy, BaseCallRatePolicy) else 1
                 policy.try_acquire(request, weight=weight)
                 return
             except CallRateLimitHit as exc:
                 last_exception = exc
+                if attempt == self._maximum_attempts_to_acquire:
+                    break
                 if block:
                     if timeout is not None:
                         time_to_wait = min(timedelta(seconds=timeout), exc.time_to_wait)

unit_tests/sources/streams/test_call_rate.py

@@ -13,6 +13,7 @@
 
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams.call_rate import (
+    AbstractCallRatePolicy,
     APIBudget,
     CallRateLimitHit,
     FixedWindowCallRatePolicy,
@@ -214,6 +215,57 @@ def test_update(self):
         policy.update(available_calls=10, call_reset_ts=None)
 
 
+class TestAPIBudgetMaximumAttempts:
+    @staticmethod
+    def _exhausted_policy() -> FixedWindowCallRatePolicy:
+        policy = FixedWindowCallRatePolicy(
+            matchers=[],
+            next_reset_ts=datetime.now() + timedelta(hours=1),
+            period=timedelta(hours=1),
+            call_limit=1,
+        )
+        policy.try_acquire(Request("GET", "https://example.com"), weight=1)
+        return policy
+
+    @pytest.mark.parametrize("maximum_attempts_to_acquire", [0, -1])
+    def test_rejects_values_below_one(self, maximum_attempts_to_acquire):
+        with pytest.raises(ValueError, match="maximum_attempts_to_acquire must be at least 1"):
+            APIBudget(policies=[], maximum_attempts_to_acquire=maximum_attempts_to_acquire)
+
+    def test_single_attempt_non_blocking_raises_when_exhausted(self):
+        budget = APIBudget(policies=[self._exhausted_policy()], maximum_attempts_to_acquire=1)
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(Request("GET", "https://example.com"), block=False)
+
+    def test_single_attempt_blocking_raises_without_sleeping(self, mocker):
+        sleep_mock = mocker.patch("airbyte_cdk.sources.streams.call_rate.time.sleep")
+        budget = APIBudget(policies=[self._exhausted_policy()], maximum_attempts_to_acquire=1)
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(Request("GET", "https://example.com"), block=True)
+        sleep_mock.assert_not_called()
+
+    @pytest.mark.parametrize("maximum_attempts_to_acquire", [1, 2, 5])
+    def test_try_acquire_called_exactly_maximum_attempts_times(
+        self, mocker, maximum_attempts_to_acquire
+    ):
+        sleep_mock = mocker.patch("airbyte_cdk.sources.streams.call_rate.time.sleep")
+        policy = mocker.Mock(spec=AbstractCallRatePolicy)
+        policy.matches.return_value = True
+        policy.try_acquire.side_effect = CallRateLimitHit(
+            error="limit hit", item=None, weight=1, rate="1/s", time_to_wait=timedelta(seconds=1)
+        )
+        budget = APIBudget(
+            policies=[policy], maximum_attempts_to_acquire=maximum_attempts_to_acquire
+        )
+
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(Request("GET", "https://example.com"), block=True)
+
+        assert policy.try_acquire.call_count == maximum_attempts_to_acquire
+        # no sleep after the final failed attempt
+        assert sleep_mock.call_count == maximum_attempts_to_acquire - 1
+
+
 class TestFixedWindowCallRatePolicy:
     def test_limit_rate(self, mocker):
         policy = FixedWindowCallRatePolicy(