Coverage for tests/test_custom_middleware_exception.py: 100%

1from pathlib import Path 7 ctx1opqrstu

2from typing import Optional 7 ctx1opqrstu

3 

4from fastapi import APIRouter, FastAPI, File, UploadFile 7 ctx1opqrstu

5from fastapi.exceptions import HTTPException 7 ctx1opqrstu

6from fastapi.testclient import TestClient 7 ctx1opqrstu

7 

8app = FastAPI() 7 ctx1opqrstu

9 

10router = APIRouter() 7 ctx1opqrstu

11 

12 

13class ContentSizeLimitMiddleware: 7 ctx1opqrstu

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): 7 ctx1opqrstu

21 self.app = app 7 ctx1abcdefg

22 self.max_content_size = max_content_size 7 ctx1abcdefg

23 

24 def receive_wrapper(self, receive): 7 ctx1opqrstu

25 received = 0 14 ctx1ahbicjdkelfmgn

26 

27 async def inner(): 14 ctx1ahbicjdkelfmgn

28 nonlocal received 

29 message = await receive() 14 ctx1ahbicjdkelfmgn

30 if message["type"] != "http.request": 14 ctx1ahbicjdkelfmgn

31 return message # pragma: no cover 

32 

33 body_len = len(message.get("body", b"")) 14 ctx1ahbicjdkelfmgn

34 received += body_len 14 ctx1ahbicjdkelfmgn

35 if received > self.max_content_size: 14 ctx1ahbicjdkelfmgn

36 raise HTTPException( 7 ctx1abcdefg

37 422, 

38 detail={ 

39 "name": "ContentSizeLimitExceeded", 

40 "code": 999, 

41 "message": "File limit exceeded", 

42 }, 

43 ) 

44 return message 7 ctx1hijklmn

45 

46 return inner 14 ctx1ahbicjdkelfmgn

47 

48 async def __call__(self, scope, receive, send): 7 ctx1opqrstu

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

50 await self.app(scope, receive, send) 14 ctx1ahbicjdkelfmgn

51 return 14 ctx1ahbicjdkelfmgn

52 

53 wrapper = self.receive_wrapper(receive) 14 ctx1ahbicjdkelfmgn

54 await self.app(scope, wrapper, send) 14 ctx1ahbicjdkelfmgn

55 

56 

57@router.post("/middleware") 7 ctx1opqrstu

58def run_middleware(file: UploadFile = File(..., description="Big File")): 7 ctx1opqrstu

59 return {"message": "OK"} 7 ctx1hijklmn

60 

61 

62app.include_router(router) 7 ctx1opqrstu

63app.add_middleware(ContentSizeLimitMiddleware, max_content_size=2**8) 7 ctx1opqrstu

64 

65 

66client = TestClient(app) 7 ctx1opqrstu

67 

68 

69def test_custom_middleware_exception(tmp_path: Path): 7 ctx1opqrstu

70 default_pydantic_max_size = 2**16 7 ctx1abcdefg

71 path = tmp_path / "test.txt" 7 ctx1abcdefg

72 path.write_bytes(b"x" * (default_pydantic_max_size + 1)) 7 ctx1abcdefg

73 

74 with client: 7 ctx1abcdefg

75 with open(path, "rb") as file: 7 ctx1abcdefg

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

77 assert response.status_code == 422, response.text 7 ctx1abcdefg

78 assert response.json() == { 7 ctx1abcdefg

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): 7 ctx1opqrstu

88 path = tmp_path / "test.txt" 7 ctx1hijklmn

89 path.write_bytes(b"<file content>") 7 ctx1hijklmn

90 

91 with client: 7 ctx1hijklmn

92 with open(path, "rb") as file: 7 ctx1hijklmn

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

94 assert response.status_code == 200, response.text 7 ctx1hijklmn

95 assert response.json() == {"message": "OK"} 7 ctx1hijklmn