openai.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. from tests.integration_tests.model_runtime.__mock.openai_completion import MockCompletionsClass
  2. from tests.integration_tests.model_runtime.__mock.openai_chat import MockChatClass
  3. from tests.integration_tests.model_runtime.__mock.openai_remote import MockModelClass
  4. from tests.integration_tests.model_runtime.__mock.openai_moderation import MockModerationClass
  5. from tests.integration_tests.model_runtime.__mock.openai_speech2text import MockSpeech2TextClass
  6. from tests.integration_tests.model_runtime.__mock.openai_embeddings import MockEmbeddingsClass
  7. from openai.resources.completions import Completions
  8. from openai.resources.chat import Completions as ChatCompletions
  9. from openai.resources.models import Models
  10. from openai.resources.moderations import Moderations
  11. from openai.resources.audio.transcriptions import Transcriptions
  12. from openai.resources.embeddings import Embeddings
  13. # import monkeypatch
  14. from _pytest.monkeypatch import MonkeyPatch
  15. from typing import Literal, Callable, List
  16. import os
  17. import pytest
  18. def mock_openai(monkeypatch: MonkeyPatch, methods: List[Literal["completion", "chat", "remote", "moderation", "speech2text", "text_embedding"]]) -> Callable[[], None]:
  19. """
  20. mock openai module
  21. :param monkeypatch: pytest monkeypatch fixture
  22. :return: unpatch function
  23. """
  24. def unpatch() -> None:
  25. monkeypatch.undo()
  26. if "completion" in methods:
  27. monkeypatch.setattr(Completions, "create", MockCompletionsClass.completion_create)
  28. if "chat" in methods:
  29. monkeypatch.setattr(ChatCompletions, "create", MockChatClass.chat_create)
  30. if "remote" in methods:
  31. monkeypatch.setattr(Models, "list", MockModelClass.list)
  32. if "moderation" in methods:
  33. monkeypatch.setattr(Moderations, "create", MockModerationClass.moderation_create)
  34. if "speech2text" in methods:
  35. monkeypatch.setattr(Transcriptions, "create", MockSpeech2TextClass.speech2text_create)
  36. if "text_embedding" in methods:
  37. monkeypatch.setattr(Embeddings, "create", MockEmbeddingsClass.create_embeddings)
  38. return unpatch
  39. MOCK = os.getenv('MOCK_SWITCH', 'false').lower() == 'true'
  40. @pytest.fixture
  41. def setup_openai_mock(request, monkeypatch):
  42. methods = request.param if hasattr(request, 'param') else []
  43. if MOCK:
  44. unpatch = mock_openai(monkeypatch, methods=methods)
  45. yield
  46. if MOCK:
  47. unpatch()