Underground-Digital/Workflow-Engine

1import os2from json import dumps3from typing import Literal4 5import httpx6import pytest7from _pytest.monkeypatch import MonkeyPatch8 9from core.helper import ssrf_proxy10 11MOCK = os.getenv("MOCK_SWITCH", "false") == "true"12 13 14class MockedHttp:15    @staticmethod16    def httpx_request(17        method: Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD"], url: str, **kwargs18    ) -> httpx.Response:19        """20        Mocked httpx.request21        """22        if url == "http://404.com":23            response = httpx.Response(status_code=404, request=httpx.Request(method, url), content=b"Not Found")24            return response25 26        # get data, files27        data = kwargs.get("data")28        files = kwargs.get("files")29        json = kwargs.get("json")30        content = kwargs.get("content")31        if data is not None:32            resp = dumps(data).encode("utf-8")33        elif files is not None:34            resp = dumps(files).encode("utf-8")35        elif json is not None:36            resp = dumps(json).encode("utf-8")37        elif content is not None:38            resp = content39        else:40            resp = b"OK"41 42        response = httpx.Response(43            status_code=200, request=httpx.Request(method, url), headers=kwargs.get("headers", {}), content=resp44        )45        return response46 47 48@pytest.fixture49def setup_http_mock(request, monkeypatch: MonkeyPatch):50    if not MOCK:51        yield52        return53 54    monkeypatch.setattr(ssrf_proxy, "make_request", MockedHttp.httpx_request)55    yield56    monkeypatch.undo()57