"""FastAPI 适配器（ASGI 纯异步版本）的端到端测试。
使用 starlette TestClient + patch OLC.async_process，
不依赖真实 Redis 或 FastAPI 模板。
"""
import unittest
from unittest.mock import patch
from fastapi import FastAPI
from fastapi.testclient import TestClient
from olc.adapters.fastapi import OlcAdapterConfig, OlcFastAPIAdapter
from olc.bean.olc_control_result import OlcControlResult
from olc.control.context.context import OlcContext
def _make_app(allow: bool):
"""构造一个最小 FastAPI 应用，OLC.async_process 被 mock。"""
app = FastAPI()
@app.get("/ping")
async def ping():
return {"ok": True}
async def fake_process(req):
ctx = OlcContext(req)
ctx.result = OlcControlResult()
ctx.result.block = not allow
return ctx
async def fake_clean(ctx):
return
patcher_p = patch("olc.adapters.fastapi.adapter.OLC.async_process", side_effect=fake_process)
patcher_c = patch("olc.adapters.fastapi.adapter.OLC.async_clean", side_effect=fake_clean)
patcher_p.start()
patcher_c.start()
app.state._patchers = [patcher_p, patcher_c]
app.add_middleware(OlcFastAPIAdapter, config=OlcAdapterConfig(token_number=1))
return app
class TestFastAPIAdapterASGI(unittest.TestCase):
def test_allowed_request_passes_through(self):
app = _make_app(allow=True)
try:
client = TestClient(app)
r = client.get("/ping")
self.assertEqual(r.status_code, 200)
self.assertEqual(r.json(), {"ok": True})
finally:
for p in app.state._patchers:
p.stop()
def test_blocked_request_returns_429(self):
app = _make_app(allow=False)
try:
client = TestClient(app)
r = client.get("/ping")
self.assertEqual(r.status_code, 429)
body = r.json()
self.assertEqual(body["error"], "Too Many Requests")
self.assertEqual(body["path"], "/ping")
finally:
for p in app.state._patchers:
p.stop()
class TestBodySizeExtraction(unittest.TestCase):
"""测试请求体大小提取"""
def test_default_extractor_reads_content_length(self):
"""默认提取器从 Content-Length 头读取"""
scope = {
"type": "http",
"headers": [(b"content-length", b"1024")],
}
result = OlcFastAPIAdapter._default_body_size_extractor_from_scope(scope)
self.assertEqual(result, 1024)
def test_default_extractor_missing_header_returns_negative_one(self):
"""Content-Length 头缺失时返回 -1"""
scope = {
"type": "http",
"headers": [],
}
result = OlcFastAPIAdapter._default_body_size_extractor_from_scope(scope)
self.assertEqual(result, -1)
def test_default_extractor_invalid_value_returns_negative_one(self):
"""Content-Length 值非整数时返回 -1"""
scope = {
"type": "http",
"headers": [(b"content-length", b"abc")],
}
result = OlcFastAPIAdapter._default_body_size_extractor_from_scope(scope)
self.assertEqual(result, -1)
def test_default_extractor_case_insensitive(self):
"""头名大小写不敏感"""
scope = {
"type": "http",
"headers": [(b"Content-Length", b"2048")],
}
result = OlcFastAPIAdapter._default_body_size_extractor_from_scope(scope)
self.assertEqual(result, 2048)
def test_default_extractor_zero_value(self):
"""Content-Length 为 0 时返回 0"""
scope = {
"type": "http",
"headers": [(b"content-length", b"0")],
}
result = OlcFastAPIAdapter._default_body_size_extractor_from_scope(scope)
self.assertEqual(result, 0)
def test_default_extractor_no_headers_key(self):
"""scope 无 headers 键时返回 -1"""
scope = {"type": "http"}
result = OlcFastAPIAdapter._default_body_size_extractor_from_scope(scope)
self.assertEqual(result, -1)
class TestBodySizeExtractorCustom(unittest.TestCase):
"""测试自定义 body_size_extractor"""
def test_custom_sync_extractor(self):
"""自定义同步提取器"""
from olc.adapters.fastapi import OlcAdapterConfig
def custom_extractor(request):
return 4096
config = OlcAdapterConfig(body_size_extractor=custom_extractor)
adapter = OlcFastAPIAdapter(app=lambda **kw: None, config=config)
scope = {"type": "http", "headers": [], "path": "/", "method": "GET"}
async def receive():
return {"type": "http.request", "body": b""}
import asyncio
result = asyncio.new_event_loop().run_until_complete(
adapter._extract_request_body_size(scope, receive)
)
self.assertEqual(result, 4096)
def test_custom_async_extractor(self):
"""自定义异步提取器"""
from olc.adapters.fastapi import OlcAdapterConfig
async def custom_extractor(request):
return 8192
config = OlcAdapterConfig(body_size_extractor=custom_extractor)
adapter = OlcFastAPIAdapter(app=lambda **kw: None, config=config)
scope = {"type": "http", "headers": [], "path": "/", "method": "GET"}
async def receive():
return {"type": "http.request", "body": b""}
import asyncio
result = asyncio.new_event_loop().run_until_complete(
adapter._extract_request_body_size(scope, receive)
)
self.assertEqual(result, 8192)
def test_custom_extractor_exception_returns_negative_one(self):
"""自定义提取器抛异常时返回 -1"""
from olc.adapters.fastapi import OlcAdapterConfig
def custom_extractor(request):
raise ValueError("test error")
config = OlcAdapterConfig(body_size_extractor=custom_extractor)
adapter = OlcFastAPIAdapter(app=lambda **kw: None, config=config)
scope = {"type": "http", "headers": [], "path": "/", "method": "GET"}
async def receive():
return {"type": "http.request", "body": b""}
import asyncio
result = asyncio.new_event_loop().run_until_complete(
adapter._extract_request_body_size(scope, receive)
)
self.assertEqual(result, -1)
if __name__ == "__main__":
unittest.main(verbosity=2)