dify/api/tests/integration_tests/workflow/nodes/__mock/http.py

import os
from json import dumps
from typing import Literal

import httpx
import pytest
from _pytest.monkeypatch import MonkeyPatch

MOCK = os.getenv('MOCK_SWITCH', 'false') == 'true'


class MockedHttp:
    def httpx_request(method: Literal['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD'],
                      url: str, **kwargs) -> httpx.Response:
        """
        Mocked httpx.request
        """
        if url == 'http://404.com':
            response = httpx.Response(
                status_code=404,
                request=httpx.Request(method, url),
                content=b'Not Found'
            )
            return response

        # get data, files
        data = kwargs.get('data', None)
        files = kwargs.get('files', None)
        if data is not None:
            resp = dumps(data).encode('utf-8')
        elif files is not None:
            resp = dumps(files).encode('utf-8')
        else:
            resp = b'OK'

        response = httpx.Response(
            status_code=200,
            request=httpx.Request(method, url),
            headers=kwargs.get('headers', {}),
            content=resp
        )
        return response


@pytest.fixture
def setup_http_mock(request, monkeypatch: MonkeyPatch):
    if not MOCK:
        yield
        return

    monkeypatch.setattr(httpx, "request", MockedHttp.httpx_request)
    yield
    monkeypatch.undo()