|
1 | | -from typing import Generator |
2 | | - |
3 | 1 | import pytest |
4 | | -from fastapi.testclient import TestClient |
5 | | -from slowapi.extension import Limiter |
6 | | -from slowapi.util import get_remote_address |
7 | | - |
8 | | -from fides.api.main import app |
9 | | -from fides.common.api.v1.urn_registry import HEALTH |
10 | | -from fides.config import CONFIG, SecuritySettings |
11 | | - |
12 | | -LIMIT = 2 |
13 | | - |
14 | | - |
15 | | -@pytest.fixture(scope="function") |
16 | | -def api_client_for_rate_limiting() -> Generator: |
17 | | - """ |
18 | | - Return a client used to make API requests ratelimited at 2/minute. |
19 | | - """ |
20 | | - app.state.limiter = Limiter( |
21 | | - default_limits=[f"{LIMIT}/minute"], |
22 | | - headers_enabled=True, |
23 | | - key_prefix=CONFIG.security.rate_limit_prefix, |
24 | | - key_func=get_remote_address, |
25 | | - retry_after="http-date", |
26 | | - ) |
27 | | - with TestClient(app) as c: |
28 | | - yield c |
29 | | - app.state.limiter = Limiter( |
30 | | - default_limits=[CONFIG.security.request_rate_limit], |
31 | | - headers_enabled=True, |
32 | | - key_prefix=CONFIG.security.rate_limit_prefix, |
33 | | - key_func=get_remote_address, |
34 | | - retry_after="http-date", |
35 | | - ) |
36 | | - |
37 | | - |
38 | | -def test_requests_rate_limited(api_client_for_rate_limiting): |
39 | | - """ |
40 | | - Asserts that incremental HTTP requests above the ratelimit threshold are |
41 | | - rebuffed from the API with a 429 response. |
42 | | -
|
43 | | - A theoretical failure condition exists in this test should the container |
44 | | - running it not be able to execute 100 requests against the client in a |
45 | | - one minute period. |
46 | | - """ |
47 | | - for _ in range(0, LIMIT): |
48 | | - response = api_client_for_rate_limiting.get(HEALTH) |
49 | | - assert response.status_code == 200 |
50 | 2 |
|
51 | | - response = api_client_for_rate_limiting.get(HEALTH) |
52 | | - assert response.status_code == 429 |
| 3 | +from fides.config import SecuritySettings |
53 | 4 |
|
54 | 5 |
|
55 | 6 | def test_rate_limit_validation(): |
|
0 commit comments