39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock
|
|
from app.middleware import (
|
|
_get_client_id_from_request, get_upload_rate_limit,
|
|
get_ai_rate_limit, upload_rate_limit_key, ai_rate_limit_key
|
|
)
|
|
|
|
def test_get_client_id_from_request_with_id():
|
|
request = MagicMock()
|
|
request.state.client_id = "test-client"
|
|
request.state.client_plan = "free"
|
|
# Key encodes client_id + plan for per-plan rate limiting
|
|
assert _get_client_id_from_request(request) == "test-client:free"
|
|
|
|
def test_get_client_id_from_request_fallback():
|
|
request = MagicMock()
|
|
del request.state.client_id # Ensure it's not there
|
|
request.client.host = "1.2.3.4"
|
|
request.headers = {}
|
|
|
|
# slowapi's get_remote_address uses request.client.host or headers
|
|
with patch("app.middleware.get_remote_address", return_value="1.2.3.4"):
|
|
assert _get_client_id_from_request(request) == "1.2.3.4"
|
|
|
|
from unittest.mock import patch
|
|
|
|
def test_rate_limit_helpers():
|
|
assert "hour" in get_upload_rate_limit("free")
|
|
assert "hour" in get_ai_rate_limit("premium")
|
|
assert get_upload_rate_limit("invalid") == get_upload_rate_limit("free")
|
|
|
|
def test_rate_limit_keys():
|
|
request = MagicMock()
|
|
request.state.client_id = "abc"
|
|
request.state.client_plan = "free"
|
|
# Keys encode client_id + plan for bucket isolation
|
|
assert upload_rate_limit_key(request) == "abc:free"
|
|
assert ai_rate_limit_key(request) == "abc:free"
|