3 changed files with 79 additions and 0 deletions
@ -0,0 +1,16 @@ |
|||||
|
import pytest |
||||
|
from fastapi.testclient import TestClient |
||||
|
from fastapi import FastAPI, Request |
||||
|
from fastapi.route_middleware import route_middleware, verify_jwt, log_route |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.post("/secure") |
||||
|
@route_middleware(verify_jwt, log_route) |
||||
|
async def secure_route(req: Request, is_true: bool): |
||||
|
|
||||
|
return {"status": "ok", "is_true": is_true,"user":req.user} |
||||
|
|
||||
|
@app.post("/open") |
||||
|
async def open_route(is_true: bool): |
||||
|
return {"status": "open", "is_true": is_true} |
||||
@ -0,0 +1,34 @@ |
|||||
|
from functools import wraps |
||||
|
from fastapi import Request |
||||
|
from typing import Callable |
||||
|
|
||||
|
def route_middleware(*middlewares:Callable): |
||||
|
def decorator(route_func:Callable): |
||||
|
@wraps(route_func) |
||||
|
async def wrapper(*args, **kwargs): |
||||
|
req = kwargs.get("req") |
||||
|
if req is None: |
||||
|
raise ValueError("Route must have 'request: Request' parameter") |
||||
|
|
||||
|
for middleware in middlewares: |
||||
|
result = middleware(req) |
||||
|
if callable(getattr(result, "__await__", None)): |
||||
|
await result |
||||
|
|
||||
|
return await route_func(*args, **kwargs) |
||||
|
|
||||
|
return wrapper |
||||
|
return decorator |
||||
|
|
||||
|
|
||||
|
|
||||
|
# Example middlewares |
||||
|
async def verify_jwt(req: Request): |
||||
|
# just a mock |
||||
|
if not (req.query_params.get("is_true") == "true"): |
||||
|
req.user={"name":"xyz","admin":True} |
||||
|
from fastapi import HTTPException, status |
||||
|
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid JWT") |
||||
|
|
||||
|
def log_route(req: Request): |
||||
|
print(f"[LOG] Path: {req.url.path}") |
||||
@ -0,0 +1,29 @@ |
|||||
|
import pytest |
||||
|
from fastapi.testclient import TestClient |
||||
|
from fastapi import FastAPI, Request |
||||
|
from fastapi.route_middleware import route_middleware, verify_jwt, log_route |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.post("/secure") |
||||
|
@route_middleware(verify_jwt, log_route) |
||||
|
async def secure_route(req: Request, is_true: bool): |
||||
|
return {"status": "ok", "is_true": is_true,"user":req.user} |
||||
|
|
||||
|
@app.post("/open") |
||||
|
async def open_route(is_true: bool): |
||||
|
return {"status": "open", "is_true": is_true} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
|
||||
|
def test_secure_route_pass(): |
||||
|
response = client.post("/secure?is_true=true") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
def test_secure_route_fail(): |
||||
|
response = client.post("/secure?is_true=false") |
||||
|
assert response.status_code == 403 |
||||
|
|
||||
|
def test_open_route(): |
||||
|
response = client.post("/open?is_true=false") |
||||
|
assert response.status_code == 200 |
||||
Loading…
Reference in new issue