Coverage for tests / test_custom_middleware_exception.py: 100%

1from pathlib import Path 4 ctx1ijkl

2 

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

4from fastapi.exceptions import HTTPException 4 ctx1ijkl

5from fastapi.testclient import TestClient 4 ctx1ijkl

6 

7app = FastAPI() 4 ctx1ijkl

8 

9router = APIRouter() 4 ctx1ijkl

10 

11 

12class ContentSizeLimitMiddleware: 4 ctx1ijkl

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 ctx1ijkl

20 self.app = app 8 ctx1abcdefgh

21 self.max_content_size = max_content_size 8 ctx1abcdefgh

22 

23 def receive_wrapper(self, receive): 4 ctx1ijkl

24 received = 0 8 ctx1abcdefgh

25 

26 async def inner(): 8 ctx1abcdefgh

27 nonlocal received 

28 message = await receive() 8 ctx1abcdefgh

29 if message["type"] != "http.request": 8 ctx1abcdefgh

30 return message # pragma: no cover 

31 

32 body_len = len(message.get("body", b"")) 8 ctx1abcdefgh

33 received += body_len 8 ctx1abcdefgh

34 if received > self.max_content_size: 8 ctx1abcdefgh

35 raise HTTPException( 4 ctx1bdfh

36 422, 

37 detail={ 

38 "name": "ContentSizeLimitExceeded", 

39 "code": 999, 

40 "message": "File limit exceeded", 

41 }, 

42 ) 

43 return message 4 ctx1aceg

44 

45 return inner 8 ctx1abcdefgh

46 

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

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

49 await self.app(scope, receive, send) 8 ctx1abcdefgh

50 return 8 ctx1abcdefgh

51 

52 wrapper = self.receive_wrapper(receive) 8 ctx1abcdefgh

53 await self.app(scope, wrapper, send) 8 ctx1abcdefgh

54 

55 

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

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

58 return {"message": "OK"} 4 ctx1aceg

59 

60 

61app.include_router(router) 4 ctx1ijkl

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

63 

64 

65client = TestClient(app) 4 ctx1ijkl

66 

67 

68def test_custom_middleware_exception(tmp_path: Path): 4 ctx1ijkl

69 default_pydantic_max_size = 2**16 4 ctx1bdfh

70 path = tmp_path / "test.txt" 4 ctx1bdfh

71 path.write_bytes(b"x" * (default_pydantic_max_size + 1)) 4 ctx1bdfh

72 

73 with client: 4 ctx1bdfh

74 with open(path, "rb") as file: 4 ctx1bdfh

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

76 assert response.status_code == 422, response.text 4 ctx1bdfh

77 assert response.json() == { 4 ctx1bdfh

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 ctx1ijkl

87 path = tmp_path / "test.txt" 4 ctx1aceg

88 path.write_bytes(b"<file content>") 4 ctx1aceg

89 

90 with client: 4 ctx1aceg

91 with open(path, "rb") as file: 4 ctx1aceg

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

93 assert response.status_code == 200, response.text 4 ctx1aceg

94 assert response.json() == {"message": "OK"} 4 ctx1aceg