Underground-Digital/Workflow-Engine

1import os2from collections.abc import Callable3from typing import Literal4 5import pytest6 7# import monkeypatch8from _pytest.monkeypatch import MonkeyPatch9from openai.resources.audio.transcriptions import Transcriptions10from openai.resources.chat import Completions as ChatCompletions11from openai.resources.completions import Completions12from openai.resources.embeddings import Embeddings13from openai.resources.models import Models14from openai.resources.moderations import Moderations15 16from tests.integration_tests.model_runtime.__mock.openai_chat import MockChatClass17from tests.integration_tests.model_runtime.__mock.openai_completion import MockCompletionsClass18from tests.integration_tests.model_runtime.__mock.openai_embeddings import MockEmbeddingsClass19from tests.integration_tests.model_runtime.__mock.openai_moderation import MockModerationClass20from tests.integration_tests.model_runtime.__mock.openai_remote import MockModelClass21from tests.integration_tests.model_runtime.__mock.openai_speech2text import MockSpeech2TextClass22 23 24def mock_openai(25    monkeypatch: MonkeyPatch,26    methods: list[Literal["completion", "chat", "remote", "moderation", "speech2text", "text_embedding"]],27) -> Callable[[], None]:28    """29    mock openai module30 31    :param monkeypatch: pytest monkeypatch fixture32    :return: unpatch function33    """34 35    def unpatch() -> None:36        monkeypatch.undo()37 38    if "completion" in methods:39        monkeypatch.setattr(Completions, "create", MockCompletionsClass.completion_create)40 41    if "chat" in methods:42        monkeypatch.setattr(ChatCompletions, "create", MockChatClass.chat_create)43 44    if "remote" in methods:45        monkeypatch.setattr(Models, "list", MockModelClass.list)46 47    if "moderation" in methods:48        monkeypatch.setattr(Moderations, "create", MockModerationClass.moderation_create)49 50    if "speech2text" in methods:51        monkeypatch.setattr(Transcriptions, "create", MockSpeech2TextClass.speech2text_create)52 53    if "text_embedding" in methods:54        monkeypatch.setattr(Embeddings, "create", MockEmbeddingsClass.create_embeddings)55 56    return unpatch57 58 59MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"60 61 62@pytest.fixture63def setup_openai_mock(request, monkeypatch):64    methods = request.param if hasattr(request, "param") else []65    if MOCK:66        unpatch = mock_openai(monkeypatch, methods=methods)67 68    yield69 70    if MOCK:71        unpatch()72