Add unit tests

Author: tolik0

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -142,6 +142,7 @@
 from airbyte_cdk.sources.declarative.transformations import AddFields, RemoveFields
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
+from airbyte_cdk.sources.streams.call_rate import MovingWindowCallRatePolicy
 from airbyte_cdk.sources.streams.concurrent.clamping import (
     ClampingEndProvider,
     DayClampingStrategy,
@@ -3564,3 +3565,82 @@ def test_create_async_retriever():
     assert isinstance(selector, RecordSelector)
     assert isinstance(extractor, DpathExtractor)
     assert extractor.field_path == ["data"]
+
+
+def test_api_budget():
+    manifest = {
+        "type": "DeclarativeSource",
+        "api_budget": {
+            "type": "HTTPAPIBudget",
+            "ratelimit_reset_header": "X-RateLimit-Reset",
+            "ratelimit_remaining_header": "X-RateLimit-Remaining",
+            "status_codes_for_ratelimit_hit": [429, 503],
+            "policies": [
+                {
+                    "type": "MovingWindowCallRatePolicy",
+                    "rates": [
+                        {
+                            "type": "Rate",
+                            "limit": 3,
+                            "interval": "PT0.1S",  # 0.1 seconds
+                        }
+                    ],
+                    "matchers": [
+                        {
+                            "type": "HttpRequestRegexMatcher",
+                            "method": "GET",
+                            "url_base": "https://api.sendgrid.com",
+                            "url_path_pattern": "/v3/marketing/lists",
+                        }
+                    ],
+                }
+            ],
+        },
+        "my_requester": {
+            "type": "HttpRequester",
+            "path": "/v3/marketing/lists",
+            "url_base": "https://api.sendgrid.com",
+            "http_method": "GET",
+            "authenticator": {
+                "type": "BasicHttpAuthenticator",
+                "username": "admin",
+                "password": "{{ config['password'] }}",
+            },
+        },
+    }
+
+    config = {
+        "password": "verysecrettoken",
+    }
+
+    factory = ModelToComponentFactory()
+    if "api_budget" in manifest:
+        factory.set_api_budget(manifest["api_budget"], config)
+
+    from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+        HttpRequester as HttpRequesterModel,
+    )
+
+    requester_definition = manifest["my_requester"]
+    assert requester_definition["type"] == "HttpRequester"
+
+    http_requester = factory.create_component(
+        model_type=HttpRequesterModel,
+        component_definition=requester_definition,
+        config=config,
+        name="lists_stream",
+        decoder=None,
+    )
+
+    assert http_requester.api_budget is not None
+    assert http_requester.api_budget.ratelimit_reset_header == "X-RateLimit-Reset"
+    assert http_requester.api_budget.status_codes_for_ratelimit_hit == [429, 503]
+    assert len(http_requester.api_budget.policies) == 1
+
+    # The single policy is a MovingWindowCallRatePolicy
+    policy = http_requester.api_budget.policies[0]
+    assert isinstance(policy, MovingWindowCallRatePolicy)
+    assert policy._bucket.rates[0].limit == 3
+    # The 0.1s from 'PT0.1S' is stored in ms by PyRateLimiter internally
+    # but here just check that the limit and interval exist
+    assert policy._bucket.rates[0].interval == 100  # 100 ms

unit_tests/sources/declarative/requesters/test_http_requester.py

@@ -2,13 +2,15 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+from datetime import timedelta
 from typing import Any, Mapping, Optional
 from unittest import mock
 from unittest.mock import MagicMock
 from urllib.parse import parse_qs, urlparse
 
 import pytest as pytest
 import requests
+import requests.sessions
 from requests import PreparedRequest
 
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
@@ -27,6 +29,12 @@
     InterpolatedRequestOptionsProvider,
 )
 from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams.call_rate import (
+    AbstractAPIBudget,
+    HttpAPIBudget,
+    MovingWindowCallRatePolicy,
+    Rate,
+)
 from airbyte_cdk.sources.streams.http.error_handlers.response_models import ResponseAction
 from airbyte_cdk.sources.streams.http.exceptions import (
     RequestBodyException,
@@ -45,6 +53,7 @@ def factory(
         request_options_provider: Optional[InterpolatedRequestOptionsProvider] = None,
         authenticator: Optional[DeclarativeAuthenticator] = None,
         error_handler: Optional[ErrorHandler] = None,
+        api_budget: Optional[HttpAPIBudget] = None,
         config: Optional[Config] = None,
         parameters: Mapping[str, Any] = None,
         disable_retries: bool = False,
@@ -61,6 +70,7 @@ def factory(
             http_method=http_method,
             request_options_provider=request_options_provider,
             error_handler=error_handler,
+            api_budget=api_budget,
             disable_retries=disable_retries,
             message_repository=message_repository or MagicMock(),
             use_cache=use_cache,
@@ -934,3 +944,25 @@ def test_backoff_strategy_from_manifest_is_respected(http_requester_factory: Any
         http_requester._http_client._request_attempt_count.get(request_mock)
         == http_requester._http_client._max_retries + 1
     )
+
+
+def test_http_requester_with_mock_apibudget(http_requester_factory, monkeypatch):
+    mock_budget = MagicMock(spec=HttpAPIBudget)
+
+    requester = http_requester_factory(
+        url_base="https://example.com",
+        path="test",
+        api_budget=mock_budget,
+    )
+
+    dummy_response = requests.Response()
+    dummy_response.status_code = 200
+    send_mock = MagicMock(return_value=dummy_response)
+    monkeypatch.setattr(requests.Session, "send", send_mock)
+
+    response = requester.send_request()
+
+    assert send_mock.call_count == 1
+    assert response.status_code == 200
+
+    assert mock_budget.acquire_call.call_count == 1