Coverage for tests/test_custom_middleware_exception.py: 100%

1from pathlib import Path 6 ctx1mnopqr

2from typing import Optional 6 ctx1mnopqr

3 

4from fastapi import APIRouter, FastAPI, File, UploadFile 6 ctx1mnopqr

5from fastapi.exceptions import HTTPException 6 ctx1mnopqr

6from fastapi.testclient import TestClient 6 ctx1mnopqr

7 

8app = FastAPI() 6 ctx1mnopqr

9 

10router = APIRouter() 6 ctx1mnopqr

11 

12 

13class ContentSizeLimitMiddleware: 6 ctx1mnopqr

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): 6 ctx1mnopqr

21 self.app = app 6 ctx1abcdef

22 self.max_content_size = max_content_size 6 ctx1abcdef

23 

24 def receive_wrapper(self, receive): 6 ctx1mnopqr

25 received = 0 12 ctx1agbhcidjekfl

26 

27 async def inner(): 12 ctx1agbhcidjekfl

28 nonlocal received 

29 message = await receive() 12 ctx1agbhcidjekfl

30 if message["type"] != "http.request": 12 ctx1agbhcidjekfl

31 return message # pragma: no cover 

32 

33 body_len = len(message.get("body", b"")) 12 ctx1agbhcidjekfl

34 received += body_len 12 ctx1agbhcidjekfl

35 if received > self.max_content_size: 12 ctx1agbhcidjekfl

36 raise HTTPException( 6 ctx1abcdef

37 422, 

38 detail={ 

39 "name": "ContentSizeLimitExceeded", 

40 "code": 999, 

41 "message": "File limit exceeded", 

42 }, 

43 ) 

44 return message 6 ctx1ghijkl

45 

46 return inner 12 ctx1agbhcidjekfl

47 

48 async def __call__(self, scope, receive, send): 6 ctx1mnopqr

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

50 await self.app(scope, receive, send) 12 ctx1agbhcidjekfl

51 return 12 ctx1agbhcidjekfl

52 

53 wrapper = self.receive_wrapper(receive) 12 ctx1agbhcidjekfl

54 await self.app(scope, wrapper, send) 12 ctx1agbhcidjekfl

55 

56 

57@router.post("/middleware") 6 ctx1mnopqr

58def run_middleware(file: UploadFile = File(..., description="Big File")): 6 ctx1mnopqr

59 return {"message": "OK"} 6 ctx1ghijkl

60 

61 

62app.include_router(router) 6 ctx1mnopqr

63app.add_middleware(ContentSizeLimitMiddleware, max_content_size=2**8) 6 ctx1mnopqr

64 

65 

66client = TestClient(app) 6 ctx1mnopqr

67 

68 

69def test_custom_middleware_exception(tmp_path: Path): 6 ctx1mnopqr

70 default_pydantic_max_size = 2**16 6 ctx1abcdef

71 path = tmp_path / "test.txt" 6 ctx1abcdef

72 path.write_bytes(b"x" * (default_pydantic_max_size + 1)) 6 ctx1abcdef

73 

74 with client: 6 ctx1abcdef

75 with open(path, "rb") as file: 6 ctx1abcdef

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

77 assert response.status_code == 422, response.text 6 ctx1abcdef

78 assert response.json() == { 6 ctx1abcdef

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): 6 ctx1mnopqr

88 path = tmp_path / "test.txt" 6 ctx1ghijkl

89 path.write_bytes(b"<file content>") 6 ctx1ghijkl

90 

91 with client: 6 ctx1ghijkl

92 with open(path, "rb") as file: 6 ctx1ghijkl

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

94 assert response.status_code == 200, response.text 6 ctx1ghijkl

95 assert response.json() == {"message": "OK"} 6 ctx1ghijkl