openai/openai-python
Publicmirrored fromhttps://github.com/openai/openai-pythonAvailable
tests/lib/snapshots.py
101lines · modecode
| 1 | from __future__ import annotations |
| 2 | |
| 3 | import os |
| 4 | import json |
| 5 | from typing import Any, Callable, Awaitable |
| 6 | from typing_extensions import TypeVar |
| 7 | |
| 8 | import httpx |
| 9 | from respx import MockRouter |
| 10 | |
| 11 | from openai import OpenAI, AsyncOpenAI |
| 12 | |
| 13 | from .utils import get_snapshot_value |
| 14 | |
| 15 | _T = TypeVar("_T") |
| 16 | |
| 17 | |
| 18 | def make_snapshot_request( |
| 19 | func: Callable[[OpenAI], _T], |
| 20 | *, |
| 21 | content_snapshot: Any, |
| 22 | respx_mock: MockRouter, |
| 23 | mock_client: OpenAI, |
| 24 | path: str, |
| 25 | ) -> _T: |
| 26 | live = os.environ.get("OPENAI_LIVE") == "1" |
| 27 | if live: |
| 28 | |
| 29 | def _on_response(response: httpx.Response) -> None: |
| 30 | # update the content snapshot |
| 31 | assert json.dumps(json.loads(response.read())) == content_snapshot |
| 32 | |
| 33 | respx_mock.stop() |
| 34 | |
| 35 | client = OpenAI( |
| 36 | http_client=httpx.Client( |
| 37 | event_hooks={ |
| 38 | "response": [_on_response], |
| 39 | } |
| 40 | ) |
| 41 | ) |
| 42 | else: |
| 43 | respx_mock.post(path).mock( |
| 44 | return_value=httpx.Response( |
| 45 | 200, |
| 46 | content=get_snapshot_value(content_snapshot), |
| 47 | headers={"content-type": "application/json"}, |
| 48 | ) |
| 49 | ) |
| 50 | |
| 51 | client = mock_client |
| 52 | |
| 53 | result = func(client) |
| 54 | |
| 55 | if live: |
| 56 | client.close() |
| 57 | |
| 58 | return result |
| 59 | |
| 60 | |
| 61 | async def make_async_snapshot_request( |
| 62 | func: Callable[[AsyncOpenAI], Awaitable[_T]], |
| 63 | *, |
| 64 | content_snapshot: Any, |
| 65 | respx_mock: MockRouter, |
| 66 | mock_client: AsyncOpenAI, |
| 67 | path: str, |
| 68 | ) -> _T: |
| 69 | live = os.environ.get("OPENAI_LIVE") == "1" |
| 70 | if live: |
| 71 | |
| 72 | async def _on_response(response: httpx.Response) -> None: |
| 73 | # update the content snapshot |
| 74 | assert json.dumps(json.loads(await response.aread())) == content_snapshot |
| 75 | |
| 76 | respx_mock.stop() |
| 77 | |
| 78 | client = AsyncOpenAI( |
| 79 | http_client=httpx.AsyncClient( |
| 80 | event_hooks={ |
| 81 | "response": [_on_response], |
| 82 | } |
| 83 | ) |
| 84 | ) |
| 85 | else: |
| 86 | respx_mock.post(path).mock( |
| 87 | return_value=httpx.Response( |
| 88 | 200, |
| 89 | content=get_snapshot_value(content_snapshot), |
| 90 | headers={"content-type": "application/json"}, |
| 91 | ) |
| 92 | ) |
| 93 | |
| 94 | client = mock_client |
| 95 | |
| 96 | result = await func(client) |
| 97 | |
| 98 | if live: |
| 99 | await client.close() |
| 100 | |
| 101 | return result |
| 102 | |