|
|
"""
|
|
|
限流器测试
|
|
|
"""
|
|
|
import pytest
|
|
|
from iti.applications import create_app
|
|
|
|
|
|
|
|
|
@pytest.fixture
|
|
|
def app():
|
|
|
"""创建测试应用"""
|
|
|
return create_app('test')
|
|
|
|
|
|
|
|
|
@pytest.fixture
|
|
|
def client(app):
|
|
|
"""创建测试客户端"""
|
|
|
return app.test_client()
|
|
|
|
|
|
|
|
|
def test_limiter_disabled_in_test_env(app):
|
|
|
"""测试环境应该禁用限流"""
|
|
|
from iti.applications.extensions.limit import limiter
|
|
|
|
|
|
# 测试环境应该禁用限流
|
|
|
assert app.config.get('RATELIMIT_ENABLED') is False # 配置中禁用
|
|
|
assert limiter is None # limiter 实例应该为 None
|
|
|
|
|
|
|
|
|
def test_limiter_config_in_dev_env():
|
|
|
"""测试开发环境限流配置"""
|
|
|
app = create_app('dev')
|
|
|
|
|
|
# 开发环境应该启用限流
|
|
|
assert app.config.get('RATELIMIT_ENABLED') is True
|
|
|
assert app.config.get('RATELIMIT_DEFAULT') == "1000 per hour"
|
|
|
assert app.config.get('RATELIMIT_STORAGE_URL') == "memory://"
|
|
|
|
|
|
|
|
|
def test_limiter_config_in_prod_env():
|
|
|
"""测试生产环境限流配置"""
|
|
|
app = create_app('prod')
|
|
|
|
|
|
# 生产环境应该启用限流
|
|
|
assert app.config.get('RATELIMIT_ENABLED') is True
|
|
|
assert app.config.get('RATELIMIT_DEFAULT') == "100 per hour"
|
|
|
# 生产环境可能使用 Redis
|
|
|
assert app.config.get('RATELIMIT_STORAGE_URL') in ["memory://", "redis://localhost:6379/0"]
|
|
|
|
|
|
|
|
|
def test_limiter_initialization():
|
|
|
"""测试限流器初始化"""
|
|
|
# 测试环境 - 应该禁用(limiter 为 None)
|
|
|
app_test = create_app('test')
|
|
|
from iti.applications.extensions.limit import limiter as test_limiter
|
|
|
assert test_limiter is None # 测试环境应该没有创建 limiter 实例
|
|
|
|
|
|
# 开发环境 - 应该启用
|
|
|
app_dev = create_app('dev')
|
|
|
from iti.applications.extensions.limit import limiter as dev_limiter
|
|
|
assert dev_limiter is not None # 开发环境应该创建了 limiter 实例
|
|
|
assert dev_limiter.enabled is True
|
|
|
|
|
|
|
|
|
def test_limiter_with_custom_config():
|
|
|
"""测试自定义限流配置"""
|
|
|
from flask import Flask
|
|
|
from iti.applications.extensions.limit import init_limiter
|
|
|
|
|
|
# 创建测试应用
|
|
|
app = Flask(__name__)
|
|
|
app.config.update({
|
|
|
'RATELIMIT_ENABLED': True,
|
|
|
'RATELIMIT_STORAGE_URL': 'memory://',
|
|
|
'RATELIMIT_DEFAULT': '50 per minute'
|
|
|
})
|
|
|
|
|
|
# 初始化限流器
|
|
|
init_limiter(app)
|
|
|
|
|
|
# 验证配置
|
|
|
from iti.applications.extensions.limit import limiter
|
|
|
assert limiter is not None
|
|
|
assert limiter.enabled is True
|
|
|
# 验证存储 URI(使用私有属性)
|
|
|
assert limiter._storage_uri == 'memory://'
|
|
|
|
|
|
|
|
|
def test_limiter_disabled_config():
|
|
|
"""测试禁用限流配置"""
|
|
|
from flask import Flask
|
|
|
from iti.applications.extensions.limit import init_limiter
|
|
|
|
|
|
# 创建测试应用
|
|
|
app = Flask(__name__)
|
|
|
app.config.update({
|
|
|
'RATELIMIT_ENABLED': False,
|
|
|
'RATELIMIT_STORAGE_URL': 'memory://',
|
|
|
'RATELIMIT_DEFAULT': '100 per hour'
|
|
|
})
|
|
|
|
|
|
# 初始化限流器
|
|
|
init_limiter(app)
|
|
|
|
|
|
# 验证限流器被禁用
|
|
|
from iti.applications.extensions.limit import limiter
|
|
|
assert limiter is None
|