Coverage for src/kwai/frontend/etag_file_response.py: 96%

1from starlette._compat import md5_hexdigest # noqa 

2from pathlib import Path 

3 

4from fastapi import Request 

5from fastapi.responses import FileResponse, Response 

6 

7 

8class EtagFileResponse(FileResponse): 

9 """A FileResponse that will check the etag when if-none-match header is passed. 

10 

11 _generate_etag uses the same implementation as FileResponse. 

12 FileResponse automatically sets the etag header with this etag value. 

13 """ 

14 

15 def __init__(self, path: Path, **kwargs) -> None: 

16 super().__init__(path, **kwargs) 

17 self._path = path 

18 

19 def _generate_etag(self) -> str: 

20 file_stat = self._path.stat() 

21 etag_base = str(file_stat.st_mtime) + "-" + str(file_stat.st_size) 

22 return f'"{md5_hexdigest(etag_base.encode(), usedforsecurity=False)}"' 

23 

24 async def __call__(self, scope, receive, send) -> None: 

25 """Check the etag, and return 304 when the file is not modified.""" 

26 request = Request(scope, receive) 

27 if_none_match = request.headers.get("if-none-match") 

28 

29 if not if_none_match: 

30 await super().__call__(scope, receive, send) 

31 return 

32 

33 etag = self._generate_etag() 

34 if if_none_match == etag: 

35 response = Response(status_code=304) 

36 await response(scope, receive, send) 

37 return 

38 

39 await super().__call__(scope, receive, send)