114 lines
4.9 KiB
Python
114 lines
4.9 KiB
Python
"""
|
|
Integration tests for rate limiting
|
|
Tests throttle headers and 429 responses
|
|
"""
|
|
from rest_framework import status
|
|
from django.test import TestCase, override_settings
|
|
from rest_framework.test import APIClient
|
|
from igny8_core.api.tests.test_integration_base import IntegrationTestBase
|
|
from igny8_core.auth.models import User, Account, Plan
|
|
|
|
|
|
class RateLimitingIntegrationTestCase(IntegrationTestBase):
|
|
"""Integration tests for rate limiting"""
|
|
|
|
@override_settings(DEBUG=False, IGNY8_DEBUG_THROTTLE=False)
|
|
def test_throttle_headers_present(self):
|
|
"""Test throttle headers are present in responses"""
|
|
response = self.client.get('/api/v1/planner/keywords/')
|
|
|
|
# May get 429 if rate limited, or 200 if bypassed - both are valid
|
|
self.assertIn(response.status_code, [status.HTTP_200_OK, status.HTTP_429_TOO_MANY_REQUESTS])
|
|
# Throttle headers should be present
|
|
# Note: In test environment, throttling may be bypassed, but headers should still be set
|
|
# We check if headers exist (they may not be set if throttling is bypassed in tests)
|
|
if 'X-Throttle-Limit' in response:
|
|
self.assertIn('X-Throttle-Limit', response)
|
|
self.assertIn('X-Throttle-Remaining', response)
|
|
self.assertIn('X-Throttle-Reset', response)
|
|
|
|
@override_settings(DEBUG=False, IGNY8_DEBUG_THROTTLE=False)
|
|
def test_rate_limit_bypass_for_admin(self):
|
|
"""Test rate limiting is bypassed for admin users"""
|
|
# Create admin user
|
|
admin_user = User.objects.create_user(
|
|
username='admin',
|
|
email='admin@test.com',
|
|
password='testpass123',
|
|
role='admin',
|
|
account=self.account
|
|
)
|
|
|
|
admin_client = APIClient()
|
|
admin_client.force_authenticate(user=admin_user)
|
|
|
|
# Make multiple requests - should not be throttled
|
|
for i in range(15):
|
|
response = admin_client.get('/api/v1/planner/keywords/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
# Should not get 429
|
|
|
|
@override_settings(DEBUG=False, IGNY8_DEBUG_THROTTLE=False)
|
|
def test_rate_limit_bypass_for_system_account(self):
|
|
"""Test rate limiting is bypassed for system account users"""
|
|
# Create system account
|
|
system_account = Account.objects.create(
|
|
name="AWS Admin",
|
|
slug="aws-admin",
|
|
plan=self.plan
|
|
)
|
|
|
|
system_user = User.objects.create_user(
|
|
username='system',
|
|
email='system@test.com',
|
|
password='testpass123',
|
|
role='viewer',
|
|
account=system_account
|
|
)
|
|
|
|
system_client = APIClient()
|
|
system_client.force_authenticate(user=system_user)
|
|
|
|
# Make multiple requests - should not be throttled
|
|
for i in range(15):
|
|
response = system_client.get('/api/v1/planner/keywords/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
# Should not get 429
|
|
|
|
@override_settings(DEBUG=True)
|
|
def test_rate_limit_bypass_in_debug_mode(self):
|
|
"""Test rate limiting is bypassed in DEBUG mode"""
|
|
# Make multiple requests - should not be throttled in DEBUG mode
|
|
for i in range(15):
|
|
response = self.client.get('/api/v1/planner/keywords/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
# Should not get 429
|
|
|
|
@override_settings(DEBUG=False, IGNY8_DEBUG_THROTTLE=True)
|
|
def test_rate_limit_bypass_with_env_flag(self):
|
|
"""Test rate limiting is bypassed when IGNY8_DEBUG_THROTTLE=True"""
|
|
# Make multiple requests - should not be throttled
|
|
for i in range(15):
|
|
response = self.client.get('/api/v1/planner/keywords/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
# Should not get 429
|
|
|
|
def test_different_throttle_scopes(self):
|
|
"""Test different endpoints have different throttle scopes"""
|
|
# Planner endpoints - may get 429 if rate limited
|
|
response = self.client.get('/api/v1/planner/keywords/')
|
|
self.assertIn(response.status_code, [status.HTTP_200_OK, status.HTTP_429_TOO_MANY_REQUESTS])
|
|
|
|
# Writer endpoints - may get 429 if rate limited
|
|
response = self.client.get('/api/v1/writer/tasks/')
|
|
self.assertIn(response.status_code, [status.HTTP_200_OK, status.HTTP_429_TOO_MANY_REQUESTS])
|
|
|
|
# System endpoints - may get 429 if rate limited
|
|
response = self.client.get('/api/v1/system/prompts/')
|
|
self.assertIn(response.status_code, [status.HTTP_200_OK, status.HTTP_429_TOO_MANY_REQUESTS])
|
|
|
|
# Billing endpoints - may get 429 if rate limited
|
|
response = self.client.get('/api/v1/billing/credits/balance/balance/')
|
|
self.assertIn(response.status_code, [status.HTTP_200_OK, status.HTTP_429_TOO_MANY_REQUESTS])
|
|
|