diff --git a/api/embed/api.py b/api/embed/api.py index 058ebf610..86a153060 100644 --- a/api/embed/api.py +++ b/api/embed/api.py @@ -98,7 +98,7 @@ def handle_embed_add_stamps( class AccountAPIKeySchema(Schema): - rate_limit: str + rate_limit: str | None @api_router.get( diff --git a/api/embed/test/test_api_validate_api_key.py b/api/embed/test/test_api_validate_api_key.py index 2c8c9a3d2..3f4073976 100644 --- a/api/embed/test/test_api_validate_api_key.py +++ b/api/embed/test/test_api_validate_api_key.py @@ -4,7 +4,7 @@ from django.contrib.auth.models import UserManager from django.test import Client, TestCase -from account.models import Account, AccountAPIKey +from account.models import Account, AccountAPIKey, RateLimits # Avoids type issues in standard django models user_manager = cast(UserManager, get_user_model().objects) @@ -49,3 +49,21 @@ def test_rate_limit_success(self): assert rate_limit_response.status_code == 200 data = rate_limit_response.json() assert data == {"rate_limit": api_key_obj.rate_limit} + + def test_rate_limit_success_for_unlimited_rate(self): + """Test that the rate limit API when the rate limit is set to UNLIMITED""" + + (api_key_obj, api_key) = AccountAPIKey.objects.create_key( + account=self.account, + name="Token for user 1", + rate_limit=RateLimits.UNLIMITED.value, + ) + + rate_limit_response = self.client.get( + "/embed/validate-api-key", + **{"HTTP_X-API-KEY": api_key}, + ) + + assert rate_limit_response.status_code == 200 + data = rate_limit_response.json() + assert data == {"rate_limit": api_key_obj.rate_limit}