Coverage for src\zapy\api\v1\api_requests.py: 82%

35 statements  

« prev     ^ index     » next       coverage.py v7.3.4, created at 2023-12-20 14:17 -0500

1from typing import Annotated 

2 

3from pydantic import BaseModel 

4from fastapi import APIRouter, Header, BackgroundTasks 

5 

6from ..deps.socketio import SocketIO 

7 

8from zapy.requests.requester import send_request, RequesterResponse, TestResult 

9from zapy.requests.models import ZapyRequest 

10from zapy.requests.exceptions import RenderLocationException 

11 

12 

13api_request_v1 = APIRouter(tags=["v1"]) 

14 

15 

16class RequestExecResponse(BaseModel): 

17 @classmethod 

18 def from_wrapper(cls, wrapper): 

19 response = wrapper.response 

20 content = response.text 

21 return cls( 

22 content=content, 

23 content_type=response.headers.get('content-type'), 

24 headers={k:v for k, v in response.headers.items()}, 

25 status=response.status_code, 

26 time=response.elapsed.total_seconds(), 

27 test_result=wrapper.test_result, 

28 ) 

29 

30 content: str 

31 content_type: str 

32 headers: dict[str,str] 

33 status: int 

34 time: float 

35 test_result: list[TestResult] 

36 

37 

38@api_request_v1.post("/request/exec") 

39async def exec_cell( 

40 sio: SocketIO, 

41 background_tasks: BackgroundTasks, 

42 zapy_request: ZapyRequest, 

43 x_request_id: Annotated[str | None, Header()] = None 

44) -> RequestExecResponse: 

45 def logger(*msg, sep=' ', end='\n'): 

46 log_message = sep.join(str(m) for m in msg) + end 

47 background_tasks.add_task(sio.emit, f'log:{x_request_id}', log_message) 

48 try: 

49 response_wrapper = await send_request( 

50 zapy_request = zapy_request, 

51 logger = logger, 

52 ) 

53 response_dict = RequestExecResponse.from_wrapper(response_wrapper) 

54 return response_dict 

55 except RenderLocationException as ex: 

56 response = ex.context.get("response") 

57 if response: 

58 response_wrapper = RequesterResponse(response) 

59 ex.context["response"] = RequestExecResponse.from_wrapper(response_wrapper) 

60 raise