Coverage for tests/test_custom_middleware_exception.py: 100%

1from pathlib import Path 5 ctx1klmno

2from typing import Optional 5 ctx1klmno

3 

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

5from fastapi.exceptions import HTTPException 5 ctx1klmno

6from fastapi.testclient import TestClient 5 ctx1klmno

7 

8app = FastAPI() 5 ctx1klmno

9 

10router = APIRouter() 5 ctx1klmno

11 

12 

13class ContentSizeLimitMiddleware: 5 ctx1klmno

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 ctx1klmno

21 self.app = app 5 ctx1abcde

22 self.max_content_size = max_content_size 5 ctx1abcde

23 

24 def receive_wrapper(self, receive): 5 ctx1klmno

25 received = 0 10 ctx1afbgchdiej

26 

27 async def inner(): 10 ctx1afbgchdiej

28 nonlocal received 

29 message = await receive() 10 ctx1afbgchdiej

30 if message["type"] != "http.request": 10 ctx1afbgchdiej

31 return message # pragma: no cover 

32 

33 body_len = len(message.get("body", b"")) 10 ctx1afbgchdiej

34 received += body_len 10 ctx1afbgchdiej

35 if received > self.max_content_size: 10 ctx1afbgchdiej

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 ctx1fghij

45 

46 return inner 10 ctx1afbgchdiej

47 

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

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

50 await self.app(scope, receive, send) 10 ctx1afbgchdiej

51 return 10 ctx1afbgchdiej

52 

53 wrapper = self.receive_wrapper(receive) 10 ctx1afbgchdiej

54 await self.app(scope, wrapper, send) 10 ctx1afbgchdiej

55 

56 

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

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

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

60 

61 

62app.include_router(router) 5 ctx1klmno

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

64 

65 

66client = TestClient(app) 5 ctx1klmno

67 

68 

69def test_custom_middleware_exception(tmp_path: Path): 5 ctx1klmno

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 ctx1klmno

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

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

90 

91 with client: 5 ctx1fghij

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

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

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

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