Coverage for tests / test_custom_middleware_exception.py: 100%

1from pathlib import Path 4 ctx1ghij

2from typing import Optional 4 ctx1ghij

3 

4from fastapi import APIRouter, FastAPI, File, UploadFile 4 ctx1ghij

5from fastapi.exceptions import HTTPException 4 ctx1ghij

6from fastapi.testclient import TestClient 4 ctx1ghij

7 

8app = FastAPI() 4 ctx1ghij

9 

10router = APIRouter() 4 ctx1ghij

11 

12 

13class ContentSizeLimitMiddleware: 4 ctx1ghij

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): 4 ctx1ghij

21 self.app = app 3 ctx1abc

22 self.max_content_size = max_content_size 3 ctx1abc

23 

24 def receive_wrapper(self, receive): 4 ctx1ghij

25 received = 0 6 ctx1adbecf

26 

27 async def inner(): 6 ctx1adbecf

28 nonlocal received 

29 message = await receive() 6 ctx1adbecf

30 if message["type"] != "http.request": 6 ctx1adbecf

31 return message # pragma: no cover 

32 

33 body_len = len(message.get("body", b"")) 6 ctx1adbecf

34 received += body_len 6 ctx1adbecf

35 if received > self.max_content_size: 6 ctx1adbecf

36 raise HTTPException( 3 ctx1abc

37 422, 

38 detail={ 

39 "name": "ContentSizeLimitExceeded", 

40 "code": 999, 

41 "message": "File limit exceeded", 

42 }, 

43 ) 

44 return message 3 ctx1def

45 

46 return inner 6 ctx1adbecf

47 

48 async def __call__(self, scope, receive, send): 4 ctx1ghij

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

50 await self.app(scope, receive, send) 6 ctx1adbecf

51 return 6 ctx1adbecf

52 

53 wrapper = self.receive_wrapper(receive) 6 ctx1adbecf

54 await self.app(scope, wrapper, send) 6 ctx1adbecf

55 

56 

57@router.post("/middleware") 4 ctx1ghij

58def run_middleware(file: UploadFile = File(..., description="Big File")): 4 ctx1ghij

59 return {"message": "OK"} 3 ctx1def

60 

61 

62app.include_router(router) 4 ctx1ghij

63app.add_middleware(ContentSizeLimitMiddleware, max_content_size=2**8) 4 ctx1ghij

64 

65 

66client = TestClient(app) 4 ctx1ghij

67 

68 

69def test_custom_middleware_exception(tmp_path: Path): 4 ctx1ghij

70 default_pydantic_max_size = 2**16 3 ctx1abc

71 path = tmp_path / "test.txt" 3 ctx1abc

72 path.write_bytes(b"x" * (default_pydantic_max_size + 1)) 3 ctx1abc

73 

74 with client: 3 ctx1abc

75 with open(path, "rb") as file: 3 ctx1abc

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

77 assert response.status_code == 422, response.text 3 ctx1abc

78 assert response.json() == { 3 ctx1abc

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): 4 ctx1ghij

88 path = tmp_path / "test.txt" 3 ctx1def

89 path.write_bytes(b"<file content>") 3 ctx1def

90 

91 with client: 3 ctx1def

92 with open(path, "rb") as file: 3 ctx1def

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

94 assert response.status_code == 200, response.text 3 ctx1def

95 assert response.json() == {"message": "OK"} 3 ctx1def