diff --git a/moonraker/components/http_client.py b/moonraker/components/http_client.py index f85a276..f05370e 100644 --- a/moonraker/components/http_client.py +++ b/moonraker/components/http_client.py @@ -12,8 +12,9 @@ import asyncio import pathlib import tempfile import logging +import copy from ..utils import ServerError -from tornado.escape import url_escape, url_unescape +from tornado.escape import url_unescape from tornado.httpclient import AsyncHTTPClient, HTTPRequest, HTTPError from tornado.httputil import HTTPHeaders from typing import ( @@ -267,9 +268,59 @@ class HttpClient: return dl.dest_file raise self.server.error(f"Retries exceeded for request: {url}") + def wrap_request(self, default_url: str, **kwargs) -> HttpRequestWrapper: + return HttpRequestWrapper(self, default_url, **kwargs) + def close(self): self.client.close() +class HttpRequestWrapper: + def __init__( + self, client: HttpClient, default_url: str, **kwargs + ) -> None: + self.client = client + self._last_response: Optional[HttpResponse] = None + self.default_request_args: Dict[str, Any] = { + "method": "GET", + "url": default_url, + } + self.default_request_args.update(kwargs) + self.request_args = copy.deepcopy(self.default_request_args) + self.reset() + + async def send(self, **kwargs) -> HttpResponse: + req_args = copy.deepcopy(self.request_args) + req_args.update(kwargs) + method = req_args.pop("method", self.default_request_args["method"]) + url = req_args.pop("url", self.default_request_args["url"]) + self._last_response = await self.client.request(method, url, **req_args) + return self._last_response + + def set_method(self, method: str) -> None: + self.request_args["method"] = method + + def set_url(self, url: str) -> None: + self.request_args["url"] = url + + def set_body( + self, body: Optional[Union[str, List[Any], Dict[str, Any]]] + ) -> None: + self.request_args["body"] = body + + def add_header(self, name: str, value: str) -> None: + headers = self.request_args.get("headers", {}) + headers[name] = value + self.request_args["headers"] = headers + + def set_headers(self, headers: Dict[str, str]) -> None: + self.request_args["headers"] = headers + + def reset(self) -> None: + self.request_args = copy.deepcopy(self.default_request_args) + + def last_response(self) -> Optional[HttpResponse]: + return self._last_response + class HttpResponse: def __init__(self, url: str,