Coverage for tests / test_custom_middleware_exception.py: 100%

1from pathlib import Path 4 ctx1ghij

2 

3from fastapi import APIRouter, FastAPI, File, UploadFile 4 ctx1ghij

4from fastapi.exceptions import HTTPException 4 ctx1ghij

5from fastapi.testclient import TestClient 4 ctx1ghij

6 

7app = FastAPI() 4 ctx1ghij

8 

9router = APIRouter() 4 ctx1ghij

10 

11 

12class ContentSizeLimitMiddleware: 4 ctx1ghij

13 """Content size limiting middleware for ASGI applications 

14 Args: 

15 app (ASGI application): ASGI application 

16 max_content_size (optional): the maximum content size allowed in bytes, None for no limit 

17 """ 

18 

19 def __init__(self, app: APIRouter, max_content_size: int | None = None): 4 ctx1ghij

20 self.app = app 3 ctx1abc

21 self.max_content_size = max_content_size 3 ctx1abc

22 

23 def receive_wrapper(self, receive): 4 ctx1ghij

24 received = 0 6 ctx1adbecf

25 

26 async def inner(): 6 ctx1adbecf

27 nonlocal received 

28 message = await receive() 6 ctx1adbecf

29 if message["type"] != "http.request": 6 ctx1adbecf

30 return message # pragma: no cover 

31 

32 body_len = len(message.get("body", b"")) 6 ctx1adbecf

33 received += body_len 6 ctx1adbecf

34 if received > self.max_content_size: 6 ctx1adbecf

35 raise HTTPException( 3 ctx1abc

36 422, 

37 detail={ 

38 "name": "ContentSizeLimitExceeded", 

39 "code": 999, 

40 "message": "File limit exceeded", 

41 }, 

42 ) 

43 return message 3 ctx1def

44 

45 return inner 6 ctx1adbecf

46 

47 async def __call__(self, scope, receive, send): 4 ctx1ghij

48 if scope["type"] != "http" or self.max_content_size is None: 6 ctx1adbecf

49 await self.app(scope, receive, send) 6 ctx1adbecf

50 return 6 ctx1adbecf

51 

52 wrapper = self.receive_wrapper(receive) 6 ctx1adbecf

53 await self.app(scope, wrapper, send) 6 ctx1adbecf

54 

55 

56@router.post("/middleware") 4 ctx1ghij

57def run_middleware(file: UploadFile = File(..., description="Big File")): 4 ctx1ghij

58 return {"message": "OK"} 3 ctx1def

59 

60 

61app.include_router(router) 4 ctx1ghij

62app.add_middleware(ContentSizeLimitMiddleware, max_content_size=2**8) 4 ctx1ghij

63 

64 

65client = TestClient(app) 4 ctx1ghij

66 

67 

68def test_custom_middleware_exception(tmp_path: Path): 4 ctx1ghij

69 default_pydantic_max_size = 2**16 3 ctx1abc

70 path = tmp_path / "test.txt" 3 ctx1abc

71 path.write_bytes(b"x" * (default_pydantic_max_size + 1)) 3 ctx1abc

72 

73 with client: 3 ctx1abc

74 with open(path, "rb") as file: 3 ctx1abc

75 response = client.post("/middleware", files={"file": file}) 3 ctx1abc

76 assert response.status_code == 422, response.text 3 ctx1abc

77 assert response.json() == { 3 ctx1abc

78 "detail": { 

79 "name": "ContentSizeLimitExceeded", 

80 "code": 999, 

81 "message": "File limit exceeded", 

82 } 

83 } 

84 

85 

86def test_custom_middleware_exception_not_raised(tmp_path: Path): 4 ctx1ghij

87 path = tmp_path / "test.txt" 3 ctx1def

88 path.write_bytes(b"<file content>") 3 ctx1def

89 

90 with client: 3 ctx1def

91 with open(path, "rb") as file: 3 ctx1def

92 response = client.post("/middleware", files={"file": file}) 3 ctx1def

93 assert response.status_code == 200, response.text 3 ctx1def

94 assert response.json() == {"message": "OK"} 3 ctx1def