openai/openai-python

Public

mirrored from https://github.com/openai/openai-pythonAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
v1.61.0

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

tests/lib/test_azure.py

251lines · modeblame

1ca831cfKrista Pratico1 years ago1import logging
6b01ba6aRobert Craigie1 years ago2from typing import Union, cast
3from typing_extensions import Literal, Protocol
08b8179aDavid Schnurr2 years ago4
6b01ba6aRobert Craigie1 years ago5import httpx
08b8179aDavid Schnurr2 years ago6import pytest
6b01ba6aRobert Craigie1 years ago7from respx import MockRouter
08b8179aDavid Schnurr2 years ago8
1ca831cfKrista Pratico1 years ago9from openai._utils import SensitiveHeadersFilter, is_dict
08b8179aDavid Schnurr2 years ago10from openai._models import FinalRequestOptions
11from openai.lib.azure import AzureOpenAI, AsyncAzureOpenAI
12
13Client = Union[AzureOpenAI, AsyncAzureOpenAI]
14
15
16sync_client = AzureOpenAI(
17api_version="2023-07-01",
18api_key="example API key",
19azure_endpoint="https://example-resource.azure.openai.com",
20)
21
22async_client = AsyncAzureOpenAI(
23api_version="2023-07-01",
24api_key="example API key",
25azure_endpoint="https://example-resource.azure.openai.com",
26)
27
28
6b01ba6aRobert Craigie1 years ago29class MockRequestCall(Protocol):
30request: httpx.Request
31
32
08b8179aDavid Schnurr2 years ago33@pytest.mark.parametrize("client", [sync_client, async_client])
34def test_implicit_deployment_path(client: Client) -> None:
35req = client._build_request(
36FinalRequestOptions.construct(
37method="post",
38url="/chat/completions",
39json_data={"model": "my-deployment-model"},
40)
41)
42assert (
43req.url
44== "https://example-resource.azure.openai.com/openai/deployments/my-deployment-model/chat/completions?api-version=2023-07-01"
45)
97a6895bStainless Bot2 years ago46
47
48@pytest.mark.parametrize(
49"client,method",
50[
51(sync_client, "copy"),
52(sync_client, "with_options"),
53(async_client, "copy"),
54(async_client, "with_options"),
55],
56)
57def test_client_copying(client: Client, method: Literal["copy", "with_options"]) -> None:
58if method == "copy":
59copied = client.copy()
60else:
61copied = client.with_options()
62
63assert copied._custom_query == {"api-version": "2023-07-01"}
64
65
66@pytest.mark.parametrize(
67"client",
68[sync_client, async_client],
69)
70def test_client_copying_override_options(client: Client) -> None:
71copied = client.copy(
72api_version="2022-05-01",
73)
74assert copied._custom_query == {"api-version": "2022-05-01"}
6b01ba6aRobert Craigie1 years ago75
76
77@pytest.mark.respx()
78def test_client_token_provider_refresh_sync(respx_mock: MockRouter) -> None:
79respx_mock.post(
80"https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-02-01"
81).mock(
82side_effect=[
83httpx.Response(500, json={"error": "server error"}),
84httpx.Response(200, json={"foo": "bar"}),
85]
86)
87
88counter = 0
89
90def token_provider() -> str:
91nonlocal counter
92
93counter += 1
94
95if counter == 1:
96return "first"
97
98return "second"
99
100client = AzureOpenAI(
101api_version="2024-02-01",
102azure_ad_token_provider=token_provider,
103azure_endpoint="https://example-resource.azure.openai.com",
104)
105client.chat.completions.create(messages=[], model="gpt-4")
106
107calls = cast("list[MockRequestCall]", respx_mock.calls)
108
109assert len(calls) == 2
110
111assert calls[0].request.headers.get("Authorization") == "Bearer first"
112assert calls[1].request.headers.get("Authorization") == "Bearer second"
113
114
115@pytest.mark.asyncio
116@pytest.mark.respx()
117async def test_client_token_provider_refresh_async(respx_mock: MockRouter) -> None:
118respx_mock.post(
119"https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-02-01"
120).mock(
121side_effect=[
122httpx.Response(500, json={"error": "server error"}),
123httpx.Response(200, json={"foo": "bar"}),
124]
125)
126
127counter = 0
128
129def token_provider() -> str:
130nonlocal counter
131
132counter += 1
133
134if counter == 1:
135return "first"
136
137return "second"
138
139client = AsyncAzureOpenAI(
140api_version="2024-02-01",
141azure_ad_token_provider=token_provider,
142azure_endpoint="https://example-resource.azure.openai.com",
143)
144
145await client.chat.completions.create(messages=[], model="gpt-4")
146
147calls = cast("list[MockRequestCall]", respx_mock.calls)
148
149assert len(calls) == 2
150
151assert calls[0].request.headers.get("Authorization") == "Bearer first"
152assert calls[1].request.headers.get("Authorization") == "Bearer second"
1ca831cfKrista Pratico1 years ago153
154
155class TestAzureLogging:
156
157@pytest.fixture(autouse=True)
158def logger_with_filter(self) -> logging.Logger:
159logger = logging.getLogger("openai")
160logger.setLevel(logging.DEBUG)
161logger.addFilter(SensitiveHeadersFilter())
162return logger
163
164@pytest.mark.respx()
165def test_azure_api_key_redacted(self, respx_mock: MockRouter, caplog: pytest.LogCaptureFixture) -> None:
166respx_mock.post(
167"https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"
168).mock(
169return_value=httpx.Response(200, json={"model": "gpt-4"})
170)
171
172client = AzureOpenAI(
173api_version="2024-06-01",
174api_key="example_api_key",
175azure_endpoint="https://example-resource.azure.openai.com",
176)
177
178with caplog.at_level(logging.DEBUG):
179client.chat.completions.create(messages=[], model="gpt-4")
180
181for record in caplog.records:
182if is_dict(record.args) and record.args.get("headers") and is_dict(record.args["headers"]):
183assert record.args["headers"]["api-key"] == "<redacted>"
184
185
186@pytest.mark.respx()
187def test_azure_bearer_token_redacted(self, respx_mock: MockRouter, caplog: pytest.LogCaptureFixture) -> None:
188respx_mock.post(
189"https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"
190).mock(
191return_value=httpx.Response(200, json={"model": "gpt-4"})
192)
193
194client = AzureOpenAI(
195api_version="2024-06-01",
196azure_ad_token="example_token",
197azure_endpoint="https://example-resource.azure.openai.com",
198)
199
200with caplog.at_level(logging.DEBUG):
201client.chat.completions.create(messages=[], model="gpt-4")
202
203for record in caplog.records:
204if is_dict(record.args) and record.args.get("headers") and is_dict(record.args["headers"]):
205assert record.args["headers"]["Authorization"] == "<redacted>"
206
207
208@pytest.mark.asyncio
209@pytest.mark.respx()
210async def test_azure_api_key_redacted_async(self, respx_mock: MockRouter, caplog: pytest.LogCaptureFixture) -> None:
211respx_mock.post(
212"https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"
213).mock(
214return_value=httpx.Response(200, json={"model": "gpt-4"})
215)
216
217client = AsyncAzureOpenAI(
218api_version="2024-06-01",
219api_key="example_api_key",
220azure_endpoint="https://example-resource.azure.openai.com",
221)
222
223with caplog.at_level(logging.DEBUG):
224await client.chat.completions.create(messages=[], model="gpt-4")
225
226for record in caplog.records:
227if is_dict(record.args) and record.args.get("headers") and is_dict(record.args["headers"]):
228assert record.args["headers"]["api-key"] == "<redacted>"
229
230
231@pytest.mark.asyncio
232@pytest.mark.respx()
233async def test_azure_bearer_token_redacted_async(self, respx_mock: MockRouter, caplog: pytest.LogCaptureFixture) -> None:
234respx_mock.post(
235"https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"
236).mock(
237return_value=httpx.Response(200, json={"model": "gpt-4"})
238)
239
240client = AsyncAzureOpenAI(
241api_version="2024-06-01",
242azure_ad_token="example_token",
243azure_endpoint="https://example-resource.azure.openai.com",
244)
245
246with caplog.at_level(logging.DEBUG):
247await client.chat.completions.create(messages=[], model="gpt-4")
248
249for record in caplog.records:
250if is_dict(record.args) and record.args.get("headers") and is_dict(record.args["headers"]):
251assert record.args["headers"]["Authorization"] == "<redacted>"