Coverage for tests/test_custom_middleware_exception.py: 100%

1from pathlib import Path 5 ctx1abcde

2from typing import Optional 5 ctx1abcde

3 

4from fastapi import APIRouter, FastAPI, File, UploadFile 5 ctx1abcde

5from fastapi.exceptions import HTTPException 5 ctx1abcde

6from fastapi.testclient import TestClient 5 ctx1abcde

7 

8app = FastAPI() 5 ctx1abcde

9 

10router = APIRouter() 5 ctx1abcde

11 

12 

13class ContentSizeLimitMiddleware: 5 ctx1abcde

14 """Content size limiting middleware for ASGI applications 

15 Args: 

16 app (ASGI application): ASGI application 

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

18 """ 

19 

20 def __init__(self, app: APIRouter, max_content_size: Optional[int] = None): 5 ctx1abcde

21 self.app = app 5 ctx1abcde

22 self.max_content_size = max_content_size 5 ctx1abcde

23 

24 def receive_wrapper(self, receive): 5 ctx1abcde

25 received = 0 5 ctx1abcde

26 

27 async def inner(): 5 ctx1abcde

28 nonlocal received 

29 message = await receive() 5 ctx1abcde

30 if message["type"] != "http.request": 5 ctx1abcde

31 return message # pragma: no cover 

32 

33 body_len = len(message.get("body", b"")) 5 ctx1abcde

34 received += body_len 5 ctx1abcde

35 if received > self.max_content_size: 5 ctx1abcde

36 raise HTTPException( 5 ctx1abcde

37 422, 

38 detail={ 

39 "name": "ContentSizeLimitExceeded", 

40 "code": 999, 

41 "message": "File limit exceeded", 

42 }, 

43 ) 

44 return message 5 ctx1abcde

45 

46 return inner 5 ctx1abcde

47 

48 async def __call__(self, scope, receive, send): 5 ctx1abcde

49 if scope["type"] != "http" or self.max_content_size is None: 5 ctx1abcde

50 await self.app(scope, receive, send) 5 ctx1abcde

51 return 5 ctx1abcde

52 

53 wrapper = self.receive_wrapper(receive) 5 ctx1abcde

54 await self.app(scope, wrapper, send) 5 ctx1abcde

55 

56 

57@router.post("/middleware") 5 ctx1abcde

58def run_middleware(file: UploadFile = File(..., description="Big File")): 5 ctx1abcde

59 return {"message": "OK"} 5 ctx1abcde

60 

61 

62app.include_router(router) 5 ctx1abcde

63app.add_middleware(ContentSizeLimitMiddleware, max_content_size=2**8) 5 ctx1abcde

64 

65 

66client = TestClient(app) 5 ctx1abcde

67 

68 

69def test_custom_middleware_exception(tmp_path: Path): 5 ctx1abcde

70 default_pydantic_max_size = 2**16 5 ctx1abcde

71 path = tmp_path / "test.txt" 5 ctx1abcde

72 path.write_bytes(b"x" * (default_pydantic_max_size + 1)) 5 ctx1abcde

73 

74 with client: 5 ctx1abcde

75 with open(path, "rb") as file: 5 ctx1abcde

76 response = client.post("/middleware", files={"file": file}) 5 ctx1abcde

77 assert response.status_code == 422, response.text 5 ctx1abcde

78 assert response.json() == { 5 ctx1abcde

79 "detail": { 

80 "name": "ContentSizeLimitExceeded", 

81 "code": 999, 

82 "message": "File limit exceeded", 

83 } 

84 } 

85 

86 

87def test_custom_middleware_exception_not_raised(tmp_path: Path): 5 ctx1abcde

88 path = tmp_path / "test.txt" 5 ctx1abcde

89 path.write_bytes(b"<file content>") 5 ctx1abcde

90 

91 with client: 5 ctx1abcde

92 with open(path, "rb") as file: 5 ctx1abcde

93 response = client.post("/middleware", files={"file": file}) 5 ctx1abcde

94 assert response.status_code == 200, response.text 5 ctx1abcde

95 assert response.json() == {"message": "OK"} 5 ctx1abcde