|
6 | 6 | import time |
7 | 7 | from datetime import datetime, timedelta |
8 | 8 | from typing import Any, Iterable, Mapping, Optional |
| 9 | +from unittest.mock import patch |
9 | 10 |
|
10 | 11 | import pytest |
11 | 12 | import requests |
|
16 | 17 | APIBudget, |
17 | 18 | CallRateLimitHit, |
18 | 19 | FixedWindowCallRatePolicy, |
| 20 | + HttpAPIBudget, |
19 | 21 | HttpRequestMatcher, |
20 | 22 | HttpRequestRegexMatcher, |
21 | 23 | MovingWindowCallRatePolicy, |
@@ -562,3 +564,149 @@ def test_combined_criteria(self): |
562 | 564 | assert not matcher(req_bad_path) |
563 | 565 | assert not matcher(req_bad_param) |
564 | 566 | assert not matcher(req_bad_header) |
| 567 | + |
| 568 | + |
| 569 | +def test_fixed_window_update_current_window_advances_past_multiple_periods(): |
| 570 | + """_update_current_window should advance past all elapsed periods, not just one.""" |
| 571 | + now = datetime.now() |
| 572 | + # Set next_reset_ts to 5 periods in the past |
| 573 | + period = timedelta(minutes=1) |
| 574 | + past_reset = now - (period * 5) |
| 575 | + policy = FixedWindowCallRatePolicy( |
| 576 | + next_reset_ts=past_reset, |
| 577 | + period=period, |
| 578 | + call_limit=10, |
| 579 | + matchers=[], |
| 580 | + ) |
| 581 | + # Trigger window update via try_acquire |
| 582 | + policy.try_acquire("request", weight=1) |
| 583 | + # After advancing, next_reset_ts should be in the future |
| 584 | + assert policy._next_reset_ts > now, ( |
| 585 | + "next_reset_ts should have advanced past now after multiple elapsed periods" |
| 586 | + ) |
| 587 | + |
| 588 | + |
| 589 | +def test_fixed_window_deadlock_scenario_429_without_ratelimit_reset(): |
| 590 | + """Reproduce the deadlock: 429 with no ratelimit-reset header should not cause extreme wait times. |
| 591 | +
|
| 592 | + The original bug chain: |
| 593 | + 1. FixedWindowCallRatePolicy created with next_reset_ts = now + 10 days |
| 594 | + 2. 429 arrives without ratelimit-reset header |
| 595 | + 3. available_calls set to 0, next_reset_ts unchanged |
| 596 | + 4. try_acquire raises CallRateLimitHit with time_to_wait ≈ 10 days |
| 597 | + """ |
| 598 | + now = datetime.now() |
| 599 | + period = timedelta(hours=1) |
| 600 | + policy = FixedWindowCallRatePolicy( |
| 601 | + next_reset_ts=now + period, |
| 602 | + period=period, |
| 603 | + call_limit=10, |
| 604 | + matchers=[], |
| 605 | + ) |
| 606 | + |
| 607 | + budget = HttpAPIBudget( |
| 608 | + policies=[policy], |
| 609 | + status_codes_for_ratelimit_hit=[429], |
| 610 | + ) |
| 611 | + |
| 612 | + # Simulate a 429 response without ratelimit-reset but with retry-after |
| 613 | + mock_response = requests.Response() |
| 614 | + mock_response.status_code = 429 |
| 615 | + mock_response.headers["retry-after"] = "60" |
| 616 | + # No ratelimit-reset header |
| 617 | + |
| 618 | + mock_request = Request("GET", "http://example.com/api") |
| 619 | + budget.update_from_response(mock_request, mock_response) |
| 620 | + |
| 621 | + # After update, available_calls should be 0 and reset_ts should be ~60s from now |
| 622 | + # The policy should NOT have a 10-day wait |
| 623 | + with pytest.raises(CallRateLimitHit) as exc_info: |
| 624 | + policy.try_acquire("request", weight=1) |
| 625 | + |
| 626 | + # The wait time should be roughly 1 hour (the period), not 10 days |
| 627 | + assert exc_info.value.time_to_wait < timedelta(hours=2), ( |
| 628 | + f"Wait time {exc_info.value.time_to_wait} is too large, likely the old 10-day bug" |
| 629 | + ) |
| 630 | + |
| 631 | + |
| 632 | +def test_http_api_budget_get_reset_ts_from_retry_after_header(): |
| 633 | + """get_reset_ts_from_response should fall back to retry-after when ratelimit-reset is absent.""" |
| 634 | + budget = HttpAPIBudget(policies=[]) |
| 635 | + |
| 636 | + mock_response = requests.Response() |
| 637 | + mock_response.status_code = 429 |
| 638 | + mock_response.headers["retry-after"] = "120" |
| 639 | + |
| 640 | + now = datetime.now() |
| 641 | + result = budget.get_reset_ts_from_response(mock_response) |
| 642 | + assert result is not None |
| 643 | + # Should be approximately 120 seconds from now |
| 644 | + expected = now + timedelta(seconds=120) |
| 645 | + assert abs((result - expected).total_seconds()) < 5, ( |
| 646 | + f"Expected reset_ts ~{expected}, got {result}" |
| 647 | + ) |
| 648 | + |
| 649 | + |
| 650 | +def test_http_api_budget_get_reset_ts_prefers_ratelimit_reset_over_retry_after(): |
| 651 | + """ratelimit-reset header should be preferred over retry-after.""" |
| 652 | + budget = HttpAPIBudget(policies=[]) |
| 653 | + |
| 654 | + mock_response = requests.Response() |
| 655 | + mock_response.status_code = 200 |
| 656 | + future_ts = int((datetime.now() + timedelta(hours=1)).timestamp()) |
| 657 | + mock_response.headers["ratelimit-reset"] = str(future_ts) |
| 658 | + mock_response.headers["retry-after"] = "30" |
| 659 | + |
| 660 | + result = budget.get_reset_ts_from_response(mock_response) |
| 661 | + assert result is not None |
| 662 | + # Should use ratelimit-reset (1 hour from now), not retry-after (30s) |
| 663 | + assert result > datetime.now() + timedelta(minutes=30) |
| 664 | + |
| 665 | + |
| 666 | +def test_http_api_budget_get_reset_ts_invalid_retry_after(): |
| 667 | + """Invalid retry-after header value should return None gracefully.""" |
| 668 | + budget = HttpAPIBudget(policies=[]) |
| 669 | + |
| 670 | + mock_response = requests.Response() |
| 671 | + mock_response.status_code = 429 |
| 672 | + mock_response.headers["retry-after"] = "not-a-number" |
| 673 | + |
| 674 | + result = budget.get_reset_ts_from_response(mock_response) |
| 675 | + assert result is None |
| 676 | + |
| 677 | + |
| 678 | +def test_http_api_budget_get_reset_ts_no_headers(): |
| 679 | + """No rate limit headers at all should return None.""" |
| 680 | + budget = HttpAPIBudget(policies=[]) |
| 681 | + |
| 682 | + mock_response = requests.Response() |
| 683 | + mock_response.status_code = 429 |
| 684 | + |
| 685 | + result = budget.get_reset_ts_from_response(mock_response) |
| 686 | + assert result is None |
| 687 | + |
| 688 | + |
| 689 | +def test_do_acquire_caps_sleep_duration(): |
| 690 | + """_do_acquire should cap sleep time to 600 seconds maximum.""" |
| 691 | + # Use call_limit=1 so try_acquire doesn't reject weight=1, then exhaust budget |
| 692 | + policy = FixedWindowCallRatePolicy( |
| 693 | + next_reset_ts=datetime.now() + timedelta(days=10), |
| 694 | + period=timedelta(days=10), |
| 695 | + call_limit=1, |
| 696 | + matchers=[], |
| 697 | + ) |
| 698 | + # Exhaust the budget so next call triggers CallRateLimitHit with ~10 day wait |
| 699 | + policy.try_acquire("warmup", weight=1) |
| 700 | + |
| 701 | + budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=2) |
| 702 | + |
| 703 | + with patch("airbyte_cdk.sources.streams.call_rate.time.sleep") as mock_sleep: |
| 704 | + mock_sleep.side_effect = [None, None] |
| 705 | + with pytest.raises(CallRateLimitHit): |
| 706 | + budget.acquire_call(Request("GET", "http://example.com"), block=True) |
| 707 | + |
| 708 | + # Sleep should have been called with at most 600 seconds |
| 709 | + assert mock_sleep.call_count > 0, "Expected sleep to be called at least once" |
| 710 | + for call_args in mock_sleep.call_args_list: |
| 711 | + sleep_seconds = call_args[0][0] |
| 712 | + assert sleep_seconds <= 600, f"Sleep duration {sleep_seconds}s exceeds the 600s cap" |
0 commit comments