GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/openai-python

Public

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

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
ae9ca46f22a01a1584c15b5db8854c017a91fa46

Branches

Tags

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

Clone

HTTPS

Download ZIP

openai-python/tests

tests/test_client.py

2803lines · modecode

RawDownload
Latest commit unavailable.
1# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
2
3from __future__ import annotations
4
5import gc
6import os
7import sys
8import json
9import asyncio
10import inspect
11import dataclasses
12import tracemalloc
13from typing import Any, Union, TypeVar, Callable, Iterable, Iterator, Optional, Protocol, Coroutine, cast
14from unittest import mock
15from typing_extensions import Literal, AsyncIterator, override
16
17import httpx
18import pytest
19from respx import MockRouter
20from pydantic import ValidationError
21
22from openai import OpenAI, AsyncOpenAI, APIResponseValidationError
23from openai.auth import WorkloadIdentity
24from openai._types import Omit
25from openai._utils import asyncify
26from openai._models import BaseModel, FinalRequestOptions
27from openai._streaming import Stream, AsyncStream
28from openai._exceptions import APIStatusError, APITimeoutError, APIResponseValidationError
29from openai._base_client import (
30 DEFAULT_TIMEOUT,
31 HTTPX_DEFAULT_TIMEOUT,
32 BaseClient,
33 OtherPlatform,
34 DefaultHttpxClient,
35 DefaultAsyncHttpxClient,
36 get_platform,
37 make_request_options,
38)
39
40from .utils import update_env
41
42T = TypeVar("T")
43base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
44api_key = "My API Key"
45admin_api_key = "My Admin API Key"
46
47
48def _get_params(client: BaseClient[Any, Any]) -> dict[str, str]:
49 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
50 url = httpx.URL(request.url)
51 return dict(url.params)
52
53
54def _low_retry_timeout(*_args: Any, **_kwargs: Any) -> float:
55 return 0.1
56
57
58def mirror_request_content(request: httpx.Request) -> httpx.Response:
59 return httpx.Response(200, content=request.content)
60
61
62# note: we can't use the httpx.MockTransport class as it consumes the request
63# body itself, which means we can't test that the body is read lazily
64class MockTransport(httpx.BaseTransport, httpx.AsyncBaseTransport):
65 def __init__(
66 self,
67 handler: Callable[[httpx.Request], httpx.Response]
68 | Callable[[httpx.Request], Coroutine[Any, Any, httpx.Response]],
69 ) -> None:
70 self.handler = handler
71
72 @override
73 def handle_request(
74 self,
75 request: httpx.Request,
76 ) -> httpx.Response:
77 assert not inspect.iscoroutinefunction(self.handler), "handler must not be a coroutine function"
78 assert inspect.isfunction(self.handler), "handler must be a function"
79 return self.handler(request)
80
81 @override
82 async def handle_async_request(
83 self,
84 request: httpx.Request,
85 ) -> httpx.Response:
86 assert inspect.iscoroutinefunction(self.handler), "handler must be a coroutine function"
87 return await self.handler(request)
88
89
90@dataclasses.dataclass
91class Counter:
92 value: int = 0
93
94
95def _make_sync_iterator(iterable: Iterable[T], counter: Optional[Counter] = None) -> Iterator[T]:
96 for item in iterable:
97 if counter:
98 counter.value += 1
99 yield item
100
101
102async def _make_async_iterator(iterable: Iterable[T], counter: Optional[Counter] = None) -> AsyncIterator[T]:
103 for item in iterable:
104 if counter:
105 counter.value += 1
106 yield item
107
108
109def _get_open_connections(client: OpenAI | AsyncOpenAI) -> int:
110 transport = client._client._transport
111 assert isinstance(transport, httpx.HTTPTransport) or isinstance(transport, httpx.AsyncHTTPTransport)
112
113 pool = transport._pool
114 return len(pool._requests)
115
116
117class TestOpenAI:
118 @pytest.mark.respx(base_url=base_url)
119 def test_raw_response(self, respx_mock: MockRouter, client: OpenAI) -> None:
120 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
121
122 response = client.post("/foo", cast_to=httpx.Response)
123 assert response.status_code == 200
124 assert isinstance(response, httpx.Response)
125 assert response.json() == {"foo": "bar"}
126
127 @pytest.mark.respx(base_url=base_url)
128 def test_raw_response_for_binary(self, respx_mock: MockRouter, client: OpenAI) -> None:
129 respx_mock.post("/foo").mock(
130 return_value=httpx.Response(200, headers={"Content-Type": "application/binary"}, content='{"foo": "bar"}')
131 )
132
133 response = client.post("/foo", cast_to=httpx.Response)
134 assert response.status_code == 200
135 assert isinstance(response, httpx.Response)
136 assert response.json() == {"foo": "bar"}
137
138 def test_copy(self, client: OpenAI) -> None:
139 copied = client.copy()
140 assert id(copied) != id(client)
141
142 copied = client.copy(api_key="another My API Key")
143 assert copied.api_key == "another My API Key"
144 assert client.api_key == "My API Key"
145
146 copied = client.copy(admin_api_key="another My Admin API Key")
147 assert copied.admin_api_key == "another My Admin API Key"
148 assert client.admin_api_key == "My Admin API Key"
149
150 def test_copy_default_options(self, client: OpenAI) -> None:
151 # options that have a default are overridden correctly
152 copied = client.copy(max_retries=7)
153 assert copied.max_retries == 7
154 assert client.max_retries == 2
155
156 copied2 = copied.copy(max_retries=6)
157 assert copied2.max_retries == 6
158 assert copied.max_retries == 7
159
160 # timeout
161 assert isinstance(client.timeout, httpx.Timeout)
162 copied = client.copy(timeout=None)
163 assert copied.timeout is None
164 assert isinstance(client.timeout, httpx.Timeout)
165
166 def test_copy_default_headers(self) -> None:
167 client = OpenAI(
168 base_url=base_url,
169 api_key=api_key,
170 admin_api_key=admin_api_key,
171 _strict_response_validation=True,
172 default_headers={"X-Foo": "bar"},
173 )
174 assert client.default_headers["X-Foo"] == "bar"
175
176 # does not override the already given value when not specified
177 copied = client.copy()
178 assert copied.default_headers["X-Foo"] == "bar"
179
180 # merges already given headers
181 copied = client.copy(default_headers={"X-Bar": "stainless"})
182 assert copied.default_headers["X-Foo"] == "bar"
183 assert copied.default_headers["X-Bar"] == "stainless"
184
185 # uses new values for any already given headers
186 copied = client.copy(default_headers={"X-Foo": "stainless"})
187 assert copied.default_headers["X-Foo"] == "stainless"
188
189 # set_default_headers
190
191 # completely overrides already set values
192 copied = client.copy(set_default_headers={})
193 assert copied.default_headers.get("X-Foo") is None
194
195 copied = client.copy(set_default_headers={"X-Bar": "Robert"})
196 assert copied.default_headers["X-Bar"] == "Robert"
197
198 with pytest.raises(
199 ValueError,
200 match="`default_headers` and `set_default_headers` arguments are mutually exclusive",
201 ):
202 client.copy(set_default_headers={}, default_headers={"X-Foo": "Bar"})
203 client.close()
204
205 def test_copy_default_query(self) -> None:
206 client = OpenAI(
207 base_url=base_url,
208 api_key=api_key,
209 admin_api_key=admin_api_key,
210 _strict_response_validation=True,
211 default_query={"foo": "bar"},
212 )
213 assert _get_params(client)["foo"] == "bar"
214
215 # does not override the already given value when not specified
216 copied = client.copy()
217 assert _get_params(copied)["foo"] == "bar"
218
219 # merges already given params
220 copied = client.copy(default_query={"bar": "stainless"})
221 params = _get_params(copied)
222 assert params["foo"] == "bar"
223 assert params["bar"] == "stainless"
224
225 # uses new values for any already given headers
226 copied = client.copy(default_query={"foo": "stainless"})
227 assert _get_params(copied)["foo"] == "stainless"
228
229 # set_default_query
230
231 # completely overrides already set values
232 copied = client.copy(set_default_query={})
233 assert _get_params(copied) == {}
234
235 copied = client.copy(set_default_query={"bar": "Robert"})
236 assert _get_params(copied)["bar"] == "Robert"
237
238 with pytest.raises(
239 ValueError,
240 # TODO: update
241 match="`default_query` and `set_default_query` arguments are mutually exclusive",
242 ):
243 client.copy(set_default_query={}, default_query={"foo": "Bar"})
244
245 client.close()
246
247 def test_copy_signature(self, client: OpenAI) -> None:
248 # ensure the same parameters that can be passed to the client are defined in the `.copy()` method
249 init_signature = inspect.signature(
250 # mypy doesn't like that we access the `__init__` property.
251 client.__init__, # type: ignore[misc]
252 )
253 copy_signature = inspect.signature(client.copy)
254 exclude_params = {"transport", "proxies", "_strict_response_validation"}
255
256 for name in init_signature.parameters.keys():
257 if name in exclude_params:
258 continue
259
260 copy_param = copy_signature.parameters.get(name)
261 assert copy_param is not None, f"copy() signature is missing the {name} param"
262
263 @pytest.mark.skipif(sys.version_info >= (3, 10), reason="fails because of a memory leak that started from 3.12")
264 def test_copy_build_request(self, client: OpenAI) -> None:
265 options = FinalRequestOptions(method="get", url="/foo")
266
267 def build_request(options: FinalRequestOptions) -> None:
268 client_copy = client.copy()
269 client_copy._build_request(options)
270
271 # ensure that the machinery is warmed up before tracing starts.
272 build_request(options)
273 gc.collect()
274
275 tracemalloc.start(1000)
276
277 snapshot_before = tracemalloc.take_snapshot()
278
279 ITERATIONS = 10
280 for _ in range(ITERATIONS):
281 build_request(options)
282
283 gc.collect()
284 snapshot_after = tracemalloc.take_snapshot()
285
286 tracemalloc.stop()
287
288 def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
289 if diff.count == 0:
290 # Avoid false positives by considering only leaks (i.e. allocations that persist).
291 return
292
293 if diff.count % ITERATIONS != 0:
294 # Avoid false positives by considering only leaks that appear per iteration.
295 return
296
297 for frame in diff.traceback:
298 if any(
299 frame.filename.endswith(fragment)
300 for fragment in [
301 # to_raw_response_wrapper leaks through the @functools.wraps() decorator.
302 #
303 # removing the decorator fixes the leak for reasons we don't understand.
304 "openai/_legacy_response.py",
305 "openai/_response.py",
306 # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
307 "openai/_compat.py",
308 # Standard library leaks we don't care about.
309 "/logging/__init__.py",
310 ]
311 ):
312 return
313
314 leaks.append(diff)
315
316 leaks: list[tracemalloc.StatisticDiff] = []
317 for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
318 add_leak(leaks, diff)
319 if leaks:
320 for leak in leaks:
321 print("MEMORY LEAK:", leak)
322 for frame in leak.traceback:
323 print(frame)
324 raise AssertionError()
325
326 def test_request_timeout(self, client: OpenAI) -> None:
327 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
328 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
329 assert timeout == DEFAULT_TIMEOUT
330
331 request = client._build_request(FinalRequestOptions(method="get", url="/foo", timeout=httpx.Timeout(100.0)))
332 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
333 assert timeout == httpx.Timeout(100.0)
334
335 def test_client_timeout_option(self) -> None:
336 client = OpenAI(
337 base_url=base_url,
338 api_key=api_key,
339 admin_api_key=admin_api_key,
340 _strict_response_validation=True,
341 timeout=httpx.Timeout(0),
342 )
343
344 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
345 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
346 assert timeout == httpx.Timeout(0)
347
348 client.close()
349
350 def test_http_client_timeout_option(self) -> None:
351 # custom timeout given to the httpx client should be used
352 with httpx.Client(timeout=None) as http_client:
353 client = OpenAI(
354 base_url=base_url,
355 api_key=api_key,
356 admin_api_key=admin_api_key,
357 _strict_response_validation=True,
358 http_client=http_client,
359 )
360
361 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
362 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
363 assert timeout == httpx.Timeout(None)
364
365 client.close()
366
367 # no timeout given to the httpx client should not use the httpx default
368 with httpx.Client() as http_client:
369 client = OpenAI(
370 base_url=base_url,
371 api_key=api_key,
372 admin_api_key=admin_api_key,
373 _strict_response_validation=True,
374 http_client=http_client,
375 )
376
377 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
378 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
379 assert timeout == DEFAULT_TIMEOUT
380
381 client.close()
382
383 # explicitly passing the default timeout currently results in it being ignored
384 with httpx.Client(timeout=HTTPX_DEFAULT_TIMEOUT) as http_client:
385 client = OpenAI(
386 base_url=base_url,
387 api_key=api_key,
388 admin_api_key=admin_api_key,
389 _strict_response_validation=True,
390 http_client=http_client,
391 )
392
393 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
394 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
395 assert timeout == DEFAULT_TIMEOUT # our default
396
397 client.close()
398
399 async def test_invalid_http_client(self) -> None:
400 with pytest.raises(TypeError, match="Invalid `http_client` arg"):
401 async with httpx.AsyncClient() as http_client:
402 OpenAI(
403 base_url=base_url,
404 api_key=api_key,
405 admin_api_key=admin_api_key,
406 _strict_response_validation=True,
407 http_client=cast(Any, http_client),
408 )
409
410 def test_default_headers_option(self) -> None:
411 test_client = OpenAI(
412 base_url=base_url,
413 api_key=api_key,
414 admin_api_key=admin_api_key,
415 _strict_response_validation=True,
416 default_headers={"X-Foo": "bar"},
417 )
418 request = test_client._build_request(FinalRequestOptions(method="get", url="/foo"))
419 assert request.headers.get("x-foo") == "bar"
420 assert request.headers.get("x-stainless-lang") == "python"
421
422 test_client2 = OpenAI(
423 base_url=base_url,
424 api_key=api_key,
425 admin_api_key=admin_api_key,
426 _strict_response_validation=True,
427 default_headers={
428 "X-Foo": "stainless",
429 "X-Stainless-Lang": "my-overriding-header",
430 },
431 )
432 request = test_client2._build_request(FinalRequestOptions(method="get", url="/foo"))
433 assert request.headers.get("x-foo") == "stainless"
434 assert request.headers.get("x-stainless-lang") == "my-overriding-header"
435
436 test_client.close()
437 test_client2.close()
438
439 def test_validate_headers(self) -> None:
440 default_headers={
441 "X-Foo": "stainless",
442 "X-Stainless-Lang": "my-overriding-header",
443 },
444 )
445 request = test_client2._build_request(FinalRequestOptions(method="get", url="/foo"))
446 assert request.headers.get("x-foo") == "stainless"
447 assert request.headers.get("x-stainless-lang") == "my-overriding-header"
448
449 test_client.close()
450 test_client2.close()
451
452 def test_validate_headers(self) -> None:
453 client = OpenAI(
454 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
455 )
456 options = client._prepare_options(FinalRequestOptions(method="get", url="/foo"))
457 request = client._build_request(options)
458
459 assert request.headers.get("Authorization") == f"Bearer {api_key}"
460
461 admin_request = client._build_request(
462 FinalRequestOptions(
463 method="get",
464 url="/organization/projects",
465 security={"admin_api_key_auth": True},
466 )
467 )
468 assert admin_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
469
470 with update_env(**{"OPENAI_API_KEY": Omit()}):
471 admin_only = OpenAI(
472 base_url=base_url,
473 api_key=None,
474 admin_api_key=admin_api_key,
475 _strict_response_validation=True,
476 )
477 admin_only_request = admin_only._build_request(
478 FinalRequestOptions(
479 method="get",
480 url="/organization/projects",
481 security={"admin_api_key_auth": True},
482 )
483 )
484 assert admin_only_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
485
486 with pytest.raises(
487 TypeError,
488 match="Could not resolve authentication method",
489 ):
490 admin_only._build_request(
491 FinalRequestOptions(
492 method="post",
493 url="/responses",
494 security={"bearer_auth": True},
495 )
496 )
497
498 with update_env(
499 **{
500 "OPENAI_API_KEY": Omit(),
501 "OPENAI_ADMIN_KEY": Omit(),
502 }
503 ):
504 with pytest.raises(OpenAIError, match="Missing credentials"):
505 OpenAI(base_url=base_url, api_key=None, admin_api_key=None, _strict_response_validation=True)
506
507 def test_workload_identity_is_mutually_exclusive_with_api_key(self) -> None:
508 with pytest.raises(
509 OpenAIError,
510 match="The `api_key` and `workload_identity` arguments are mutually exclusive",
511 ):
512 OpenAI(
513 base_url=base_url,
514 api_key=api_key,
515 workload_identity=workload_identity, # type: ignore[reportArgumentType]
516 organization="org_123",
517 _strict_response_validation=True,
518 )
519
520 def test_default_query_option(self) -> None:
521 client = OpenAI(
522 base_url=base_url,
523 api_key=api_key,
524 admin_api_key=admin_api_key,
525 _strict_response_validation=True,
526 default_query={"query_param": "bar"},
527 )
528 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
529 url = httpx.URL(request.url)
530 assert dict(url.params) == {"query_param": "bar"}
531
532 request = client._build_request(
533 FinalRequestOptions(
534 method="get",
535 url="/foo",
536 params={"foo": "baz", "query_param": "overridden"},
537 )
538 )
539 url = httpx.URL(request.url)
540 assert dict(url.params) == {"foo": "baz", "query_param": "overridden"}
541
542 client.close()
543
544 def test_hardcoded_query_params_in_url(self, client: OpenAI) -> None:
545 request = client._build_request(FinalRequestOptions(method="get", url="/foo?beta=true"))
546 url = httpx.URL(request.url)
547 assert dict(url.params) == {"beta": "true"}
548
549 request = client._build_request(
550 FinalRequestOptions(
551 method="get",
552 url="/foo?beta=true",
553 params={"limit": "10", "page": "abc"},
554 )
555 )
556 url = httpx.URL(request.url)
557 assert dict(url.params) == {"beta": "true", "limit": "10", "page": "abc"}
558
559 request = client._build_request(
560 FinalRequestOptions(
561 method="get",
562 url="/files/a%2Fb?beta=true",
563 params={"limit": "10"},
564 )
565 )
566 assert request.url.raw_path == b"/files/a%2Fb?beta=true&limit=10"
567
568 def test_request_extra_json(self, client: OpenAI) -> None:
569 request = client._build_request(
570 FinalRequestOptions(
571 method="post",
572 url="/foo",
573 json_data={"foo": "bar"},
574 extra_json={"baz": False},
575 ),
576 )
577 data = json.loads(request.content.decode("utf-8"))
578 assert data == {"foo": "bar", "baz": False}
579
580 request = client._build_request(
581 FinalRequestOptions(
582 method="post",
583 url="/foo",
584 extra_json={"baz": False},
585 ),
586 )
587 data = json.loads(request.content.decode("utf-8"))
588 assert data == {"baz": False}
589
590 # `extra_json` takes priority over `json_data` when keys clash
591 request = client._build_request(
592 FinalRequestOptions(
593 method="post",
594 url="/foo",
595 json_data={"foo": "bar", "baz": True},
596 extra_json={"baz": None},
597 ),
598 )
599 data = json.loads(request.content.decode("utf-8"))
600 assert data == {"foo": "bar", "baz": None}
601
602 def test_request_extra_headers(self, client: OpenAI) -> None:
603 request = client._build_request(
604 FinalRequestOptions(
605 method="post",
606 url="/foo",
607 **make_request_options(extra_headers={"X-Foo": "Foo"}),
608 ),
609 )
610 assert request.headers.get("X-Foo") == "Foo"
611
612 # `extra_headers` takes priority over `default_headers` when keys clash
613 request = client.with_options(default_headers={"X-Bar": "true"})._build_request(
614 FinalRequestOptions(
615 method="post",
616 url="/foo",
617 **make_request_options(
618 extra_headers={"X-Bar": "false"},
619 ),
620 ),
621 )
622 assert request.headers.get("X-Bar") == "false"
623
624 def test_request_extra_query(self, client: OpenAI) -> None:
625 request = client._build_request(
626 FinalRequestOptions(
627 method="post",
628 url="/foo",
629 **make_request_options(
630 extra_query={"my_query_param": "Foo"},
631 ),
632 ),
633 )
634 params = dict(request.url.params)
635 assert params == {"my_query_param": "Foo"}
636
637 # if both `query` and `extra_query` are given, they are merged
638 request = client._build_request(
639 FinalRequestOptions(
640 method="post",
641 url="/foo",
642 **make_request_options(
643 query={"bar": "1"},
644 extra_query={"foo": "2"},
645 ),
646 ),
647 )
648 params = dict(request.url.params)
649 assert params == {"bar": "1", "foo": "2"}
650
651 # `extra_query` takes priority over `query` when keys clash
652 request = client._build_request(
653 FinalRequestOptions(
654 method="post",
655 url="/foo",
656 **make_request_options(
657 query={"foo": "1"},
658 extra_query={"foo": "2"},
659 ),
660 ),
661 )
662 params = dict(request.url.params)
663 assert params == {"foo": "2"}
664
665 def test_multipart_repeating_array(self, client: OpenAI) -> None:
666 request = client._build_request(
667 FinalRequestOptions.construct(
668 method="post",
669 url="/foo",
670 headers={"Content-Type": "multipart/form-data; boundary=6b7ba517decee4a450543ea6ae821c82"},
671 json_data={"array": ["foo", "bar"]},
672 files=[("foo.txt", b"hello world")],
673 )
674 )
675
676 assert request.read().split(b"\r\n") == [
677 b"--6b7ba517decee4a450543ea6ae821c82",
678 b'Content-Disposition: form-data; name="array[]"',
679 b"",
680 b"foo",
681 b"--6b7ba517decee4a450543ea6ae821c82",
682 b'Content-Disposition: form-data; name="array[]"',
683 b"",
684 b"bar",
685 b"--6b7ba517decee4a450543ea6ae821c82",
686 b'Content-Disposition: form-data; name="foo.txt"; filename="upload"',
687 b"Content-Type: application/octet-stream",
688 b"",
689 b"hello world",
690 b"--6b7ba517decee4a450543ea6ae821c82--",
691 b"",
692 ]
693
694 @pytest.mark.respx(base_url=base_url)
695 def test_binary_content_upload(self, respx_mock: MockRouter, client: OpenAI) -> None:
696 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
697
698 file_content = b"Hello, this is a test file."
699
700 response = client.post(
701 "/upload",
702 content=file_content,
703 cast_to=httpx.Response,
704 options={"headers": {"Content-Type": "application/octet-stream"}},
705 )
706
707 assert response.status_code == 200
708 assert response.request.headers["Content-Type"] == "application/octet-stream"
709 assert response.content == file_content
710
711 def test_binary_content_upload_with_iterator(self) -> None:
712 file_content = b"Hello, this is a test file."
713 counter = Counter()
714 iterator = _make_sync_iterator([file_content], counter=counter)
715
716 def mock_handler(request: httpx.Request) -> httpx.Response:
717 assert counter.value == 0, "the request body should not have been read"
718 return httpx.Response(200, content=request.read())
719
720 with OpenAI(
721 base_url=base_url,
722 api_key=api_key,
723 admin_api_key=admin_api_key,
724 _strict_response_validation=True,
725 http_client=httpx.Client(transport=MockTransport(handler=mock_handler)),
726 ) as client:
727 response = client.post(
728 "/upload",
729 content=iterator,
730 cast_to=httpx.Response,
731 options={"headers": {"Content-Type": "application/octet-stream"}},
732 )
733
734 assert response.status_code == 200
735 assert response.request.headers["Content-Type"] == "application/octet-stream"
736 assert response.content == file_content
737 assert counter.value == 1
738
739 @pytest.mark.respx(base_url=base_url)
740 def test_binary_content_upload_with_body_is_deprecated(self, respx_mock: MockRouter, client: OpenAI) -> None:
741 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
742
743 file_content = b"Hello, this is a test file."
744
745 with pytest.deprecated_call(
746 match="Passing raw bytes as `body` is deprecated and will be removed in a future version. Please pass raw bytes via the `content` parameter instead."
747 ):
748 response = client.post(
749 "/upload",
750 body=file_content,
751 cast_to=httpx.Response,
752 options={"headers": {"Content-Type": "application/octet-stream"}},
753 )
754
755 assert response.status_code == 200
756 assert response.request.headers["Content-Type"] == "application/octet-stream"
757 assert response.content == file_content
758
759 @pytest.mark.respx(base_url=base_url)
760 def test_basic_union_response(self, respx_mock: MockRouter, client: OpenAI) -> None:
761 class Model1(BaseModel):
762 name: str
763
764 class Model2(BaseModel):
765 foo: str
766
767 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
768
769 response = client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
770 assert isinstance(response, Model2)
771 assert response.foo == "bar"
772
773 @pytest.mark.respx(base_url=base_url)
774 def test_union_response_different_types(self, respx_mock: MockRouter, client: OpenAI) -> None:
775 """Union of objects with the same field name using a different type"""
776
777 class Model1(BaseModel):
778 foo: int
779
780 class Model2(BaseModel):
781 foo: str
782
783 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
784
785 response = client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
786 assert isinstance(response, Model2)
787 assert response.foo == "bar"
788
789 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": 1}))
790
791 response = client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
792 assert isinstance(response, Model1)
793 assert response.foo == 1
794
795 @pytest.mark.respx(base_url=base_url)
796 def test_non_application_json_content_type_for_json_data(self, respx_mock: MockRouter, client: OpenAI) -> None:
797 """
798 Response that sets Content-Type to something other than application/json but returns json data
799 """
800
801 class Model(BaseModel):
802 foo: int
803
804 respx_mock.get("/foo").mock(
805 return_value=httpx.Response(
806 200,
807 content=json.dumps({"foo": 2}),
808 headers={"Content-Type": "application/text"},
809 )
810 )
811
812 response = client.get("/foo", cast_to=Model)
813 assert isinstance(response, Model)
814 assert response.foo == 2
815
816 def test_base_url_setter(self) -> None:
817 client = OpenAI(
818 base_url="https://example.com/from_init",
819 api_key=api_key,
820 admin_api_key=admin_api_key,
821 _strict_response_validation=True,
822 )
823 assert client.base_url == "https://example.com/from_init/"
824
825 client.base_url = "https://example.com/from_setter" # type: ignore[assignment]
826
827 assert client.base_url == "https://example.com/from_setter/"
828
829 client.close()
830
831 def test_base_url_env(self) -> None:
832 with update_env(OPENAI_BASE_URL="http://localhost:5000/from/env"):
833 client = OpenAI(api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True)
834 assert client.base_url == "http://localhost:5000/from/env/"
835
836 @pytest.mark.parametrize(
837 "client",
838 [
839 OpenAI(
840 base_url="http://localhost:5000/custom/path/",
841 api_key=api_key,
842 admin_api_key=admin_api_key,
843 _strict_response_validation=True,
844 ),
845 OpenAI(
846 base_url="http://localhost:5000/custom/path/",
847 api_key=api_key,
848 admin_api_key=admin_api_key,
849 _strict_response_validation=True,
850 http_client=httpx.Client(),
851 ),
852 ],
853 ids=["standard", "custom http client"],
854 )
855 def test_base_url_trailing_slash(self, client: OpenAI) -> None:
856 request = client._build_request(
857 FinalRequestOptions(
858 method="post",
859 url="/foo",
860 json_data={"foo": "bar"},
861 ),
862 )
863 assert request.url == "http://localhost:5000/custom/path/foo"
864 client.close()
865
866 @pytest.mark.parametrize(
867 "client",
868 [
869 OpenAI(
870 base_url="http://localhost:5000/custom/path/",
871 api_key=api_key,
872 admin_api_key=admin_api_key,
873 _strict_response_validation=True,
874 ),
875 OpenAI(
876 base_url="http://localhost:5000/custom/path/",
877 api_key=api_key,
878 admin_api_key=admin_api_key,
879 _strict_response_validation=True,
880 http_client=httpx.Client(),
881 ),
882 ],
883 ids=["standard", "custom http client"],
884 )
885 def test_base_url_no_trailing_slash(self, client: OpenAI) -> None:
886 request = client._build_request(
887 FinalRequestOptions(
888 method="post",
889 url="/foo",
890 json_data={"foo": "bar"},
891 ),
892 )
893 assert request.url == "http://localhost:5000/custom/path/foo"
894 client.close()
895
896 @pytest.mark.parametrize(
897 "client",
898 [
899 OpenAI(
900 base_url="http://localhost:5000/custom/path/",
901 api_key=api_key,
902 admin_api_key=admin_api_key,
903 _strict_response_validation=True,
904 ),
905 OpenAI(
906 base_url="http://localhost:5000/custom/path/",
907 api_key=api_key,
908 admin_api_key=admin_api_key,
909 _strict_response_validation=True,
910 http_client=httpx.Client(),
911 ),
912 ],
913 ids=["standard", "custom http client"],
914 )
915 def test_absolute_request_url(self, client: OpenAI) -> None:
916 request = client._build_request(
917 FinalRequestOptions(
918 method="post",
919 url="https://myapi.com/foo",
920 json_data={"foo": "bar"},
921 ),
922 )
923 assert request.url == "https://myapi.com/foo"
924 client.close()
925
926 def test_copied_client_does_not_close_http(self) -> None:
927 test_client = OpenAI(
928 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
929 )
930 assert not test_client.is_closed()
931
932 copied = test_client.copy()
933 assert copied is not test_client
934
935 del copied
936
937 assert not test_client.is_closed()
938
939 def test_client_context_manager(self) -> None:
940 test_client = OpenAI(
941 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
942 )
943 with test_client as c2:
944 assert c2 is test_client
945 assert not c2.is_closed()
946 assert not test_client.is_closed()
947 assert test_client.is_closed()
948
949 @pytest.mark.respx(base_url=base_url)
950 def test_client_response_validation_error(self, respx_mock: MockRouter, client: OpenAI) -> None:
951 class Model(BaseModel):
952 foo: str
953
954 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": {"invalid": True}}))
955
956 with pytest.raises(APIResponseValidationError) as exc:
957 client.get("/foo", cast_to=Model)
958
959 assert isinstance(exc.value.__cause__, ValidationError)
960
961 def test_client_max_retries_validation(self) -> None:
962 with pytest.raises(TypeError, match=r"max_retries cannot be None"):
963 OpenAI(
964 base_url=base_url,
965 api_key=api_key,
966 admin_api_key=admin_api_key,
967 _strict_response_validation=True,
968 max_retries=cast(Any, None),
969 )
970
971 @pytest.mark.respx(base_url=base_url)
972 def test_default_stream_cls(self, respx_mock: MockRouter, client: OpenAI) -> None:
973 class Model(BaseModel):
974 name: str
975
976 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
977
978 stream = client.post("/foo", cast_to=Model, stream=True, stream_cls=Stream[Model])
979 assert isinstance(stream, Stream)
980 stream.response.close()
981
982 @pytest.mark.respx(base_url=base_url)
983 def test_received_text_for_expected_json(self, respx_mock: MockRouter) -> None:
984 class Model(BaseModel):
985 name: str
986
987 respx_mock.get("/foo").mock(return_value=httpx.Response(200, text="my-custom-format"))
988
989 strict_client = OpenAI(
990 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
991 )
992
993 with pytest.raises(APIResponseValidationError):
994 strict_client.get("/foo", cast_to=Model)
995
996 non_strict_client = OpenAI(
997 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=False
998 )
999
1000 response = non_strict_client.get("/foo", cast_to=Model)
1001 assert isinstance(response, str) # type: ignore[unreachable]
1002
1003 strict_client.close()
1004 non_strict_client.close()
1005
1006 @pytest.mark.parametrize(
1007 "remaining_retries,retry_after,timeout",
1008 [
1009 [3, "20", 20],
1010 [3, "0", 0.5],
1011 [3, "-10", 0.5],
1012 [3, "60", 60],
1013 [3, "61", 0.5],
1014 [3, "Fri, 29 Sep 2023 16:26:57 GMT", 20],
1015 [3, "Fri, 29 Sep 2023 16:26:37 GMT", 0.5],
1016 [3, "Fri, 29 Sep 2023 16:26:27 GMT", 0.5],
1017 [3, "Fri, 29 Sep 2023 16:27:37 GMT", 60],
1018 [3, "Fri, 29 Sep 2023 16:27:38 GMT", 0.5],
1019 [3, "99999999999999999999999999999999999", 0.5],
1020 [3, "Zun, 29 Sep 2023 16:26:27 GMT", 0.5],
1021 [3, "", 0.5],
1022 [2, "", 0.5 * 2.0],
1023 [1, "", 0.5 * 4.0],
1024 [-1100, "", 8], # test large number potentially overflowing
1025 ],
1026 )
1027 @mock.patch("time.time", mock.MagicMock(return_value=1696004797))
1028 def test_parse_retry_after_header(
1029 self, remaining_retries: int, retry_after: str, timeout: float, client: OpenAI
1030 ) -> None:
1031 headers = httpx.Headers({"retry-after": retry_after})
1032 options = FinalRequestOptions(method="get", url="/foo", max_retries=3)
1033 calculated = client._calculate_retry_timeout(remaining_retries, options, headers)
1034 assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType]
1035
1036 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1037 @pytest.mark.respx(base_url=base_url)
1038 def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: OpenAI) -> None:
1039 respx_mock.post("/chat/completions").mock(side_effect=httpx.TimeoutException("Test timeout error"))
1040
1041 with pytest.raises(APITimeoutError):
1042 client.chat.completions.with_streaming_response.create(
1043 messages=[
1044 {
1045 "content": "string",
1046 "role": "developer",
1047 }
1048 ],
1049 model="gpt-5.4",
1050 ).__enter__()
1051
1052 assert _get_open_connections(client) == 0
1053
1054 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1055 @pytest.mark.respx(base_url=base_url)
1056 def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: OpenAI) -> None:
1057 respx_mock.post("/chat/completions").mock(return_value=httpx.Response(500))
1058
1059 with pytest.raises(APIStatusError):
1060 client.chat.completions.with_streaming_response.create(
1061 messages=[
1062 {
1063 "content": "string",
1064 "role": "developer",
1065 }
1066 ],
1067 model="gpt-5.4",
1068 ).__enter__()
1069 assert _get_open_connections(client) == 0
1070
1071 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1072 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1073 @pytest.mark.respx(base_url=base_url)
1074 @pytest.mark.parametrize("failure_mode", ["status", "exception"])
1075 def test_retries_taken(
1076 self,
1077 client: OpenAI,
1078 failures_before_success: int,
1079 failure_mode: Literal["status", "exception"],
1080 respx_mock: MockRouter,
1081 ) -> None:
1082 client = client.with_options(max_retries=4)
1083
1084 nb_retries = 0
1085
1086 def retry_handler(_request: httpx.Request) -> httpx.Response:
1087 nonlocal nb_retries
1088 if nb_retries < failures_before_success:
1089 nb_retries += 1
1090 if failure_mode == "exception":
1091 raise RuntimeError("oops")
1092 return httpx.Response(500)
1093 return httpx.Response(200)
1094
1095 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1096
1097 response = client.chat.completions.with_raw_response.create(
1098 messages=[
1099 {
1100 "content": "string",
1101 "role": "developer",
1102 }
1103 ],
1104 model="gpt-5.4",
1105 )
1106
1107 assert response.retries_taken == failures_before_success
1108 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
1109
1110 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1111 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1112 @pytest.mark.respx(base_url=base_url)
1113 def test_omit_retry_count_header(
1114 self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter
1115 ) -> None:
1116 client = client.with_options(max_retries=4)
1117
1118 nb_retries = 0
1119
1120 def retry_handler(_request: httpx.Request) -> httpx.Response:
1121 nonlocal nb_retries
1122 if nb_retries < failures_before_success:
1123 nb_retries += 1
1124 return httpx.Response(500)
1125 return httpx.Response(200)
1126
1127 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1128
1129 response = client.chat.completions.with_raw_response.create(
1130 messages=[
1131 {
1132 "content": "string",
1133 "role": "developer",
1134 }
1135 ],
1136 model="gpt-5.4",
1137 extra_headers={"x-stainless-retry-count": Omit()},
1138 )
1139
1140 assert len(response.http_request.headers.get_list("x-stainless-retry-count")) == 0
1141
1142 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1143 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1144 @pytest.mark.respx(base_url=base_url)
1145 def test_overwrite_retry_count_header(
1146 self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter
1147 ) -> None:
1148 client = client.with_options(max_retries=4)
1149
1150 nb_retries = 0
1151
1152 def retry_handler(_request: httpx.Request) -> httpx.Response:
1153 nonlocal nb_retries
1154 if nb_retries < failures_before_success:
1155 nb_retries += 1
1156 return httpx.Response(500)
1157 return httpx.Response(200)
1158
1159 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1160
1161 response = client.chat.completions.with_raw_response.create(
1162 messages=[
1163 {
1164 "content": "string",
1165 "role": "developer",
1166 }
1167 ],
1168 model="gpt-5.4",
1169 extra_headers={"x-stainless-retry-count": "42"},
1170 )
1171
1172 assert response.http_request.headers.get("x-stainless-retry-count") == "42"
1173
1174 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1175 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1176 @pytest.mark.respx(base_url=base_url)
1177 def test_retries_taken_new_response_class(
1178 self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter
1179 ) -> None:
1180 client = client.with_options(max_retries=4)
1181
1182 nb_retries = 0
1183
1184 def retry_handler(_request: httpx.Request) -> httpx.Response:
1185 nonlocal nb_retries
1186 if nb_retries < failures_before_success:
1187 nb_retries += 1
1188 return httpx.Response(500)
1189 return httpx.Response(200)
1190
1191 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1192
1193 with client.chat.completions.with_streaming_response.create(
1194 messages=[
1195 {
1196 "content": "string",
1197 "role": "developer",
1198 }
1199 ],
1200 model="gpt-5.4",
1201 ) as response:
1202 assert response.retries_taken == failures_before_success
1203 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
1204
1205 def test_proxy_environment_variables(self, monkeypatch: pytest.MonkeyPatch) -> None:
1206 # Test that the proxy environment variables are set correctly
1207 monkeypatch.setenv("HTTPS_PROXY", "https://example.org")
1208 # Delete in case our environment has any proxy env vars set
1209 monkeypatch.delenv("HTTP_PROXY", raising=False)
1210 monkeypatch.delenv("ALL_PROXY", raising=False)
1211 monkeypatch.delenv("NO_PROXY", raising=False)
1212 monkeypatch.delenv("http_proxy", raising=False)
1213 monkeypatch.delenv("https_proxy", raising=False)
1214 monkeypatch.delenv("all_proxy", raising=False)
1215 monkeypatch.delenv("no_proxy", raising=False)
1216
1217 client = DefaultHttpxClient()
1218
1219 mounts = tuple(client._mounts.items())
1220 assert len(mounts) == 1
1221 assert mounts[0][0].pattern == "https://"
1222
1223 @pytest.mark.filterwarnings("ignore:.*deprecated.*:DeprecationWarning")
1224 def test_default_client_creation(self) -> None:
1225 # Ensure that the client can be initialized without any exceptions
1226 DefaultHttpxClient(
1227 verify=True,
1228 cert=None,
1229 trust_env=True,
1230 http1=True,
1231 http2=False,
1232 limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
1233 )
1234
1235 @pytest.mark.respx(base_url=base_url)
1236 def test_follow_redirects(self, respx_mock: MockRouter, client: OpenAI) -> None:
1237 # Test that the default follow_redirects=True allows following redirects
1238 respx_mock.post("/redirect").mock(
1239 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1240 )
1241 respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
1242
1243 response = client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
1244 assert response.status_code == 200
1245 assert response.json() == {"status": "ok"}
1246
1247 @pytest.mark.respx(base_url=base_url)
1248 def test_follow_redirects_disabled(self, respx_mock: MockRouter, client: OpenAI) -> None:
1249 # Test that follow_redirects=False prevents following redirects
1250 respx_mock.post("/redirect").mock(
1251 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1252 )
1253
1254 with pytest.raises(APIStatusError) as exc_info:
1255 client.post("/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response)
1256
1257 assert exc_info.value.response.status_code == 302
1258 assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
1259
1260 def test_api_key_before_after_refresh_provider(self) -> None:
1261 client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token")
1262
1263 assert client.api_key == ""
1264 assert "Authorization" not in client.auth_headers
1265
1266 client._refresh_api_key()
1267
1268 assert client.api_key == "test_bearer_token"
1269 assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
1270
1271 def test_api_key_before_after_refresh_str(self) -> None:
1272 client = OpenAI(base_url=base_url, api_key="test_api_key")
1273
1274 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
1275 client._refresh_api_key()
1276
1277 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
1278
1279 @pytest.mark.respx()
1280 def test_api_key_refresh_on_retry(self, respx_mock: MockRouter) -> None:
1281 respx_mock.post(base_url + "/chat/completions").mock(
1282 side_effect=[
1283 httpx.Response(500, json={"error": "server error"}),
1284 httpx.Response(200, json={"foo": "bar"}),
1285 ]
1286 )
1287
1288 counter = 0
1289
1290 def token_provider() -> str:
1291 nonlocal counter
1292
1293 counter += 1
1294
1295 if counter == 1:
1296 return "first"
1297
1298 return "second"
1299
1300 client = OpenAI(base_url=base_url, api_key=token_provider)
1301 client.chat.completions.create(messages=[], model="gpt-4")
1302
1303 calls = cast("list[MockRequestCall]", respx_mock.calls)
1304 assert len(calls) == 2
1305
1306 assert calls[0].request.headers.get("Authorization") == "Bearer first"
1307 assert calls[1].request.headers.get("Authorization") == "Bearer second"
1308
1309 def test_copy_auth(self) -> None:
1310 client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token_1").copy(
1311 api_key=lambda: "test_bearer_token_2"
1312 )
1313 client._refresh_api_key()
1314 assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}
1315
1316
1317class TestAsyncOpenAI:
1318 @pytest.mark.respx(base_url=base_url)
1319 async def test_raw_response(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
1320 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
1321
1322 response = await async_client.post("/foo", cast_to=httpx.Response)
1323 assert response.status_code == 200
1324 assert isinstance(response, httpx.Response)
1325 assert response.json() == {"foo": "bar"}
1326
1327 @pytest.mark.respx(base_url=base_url)
1328 async def test_raw_response_for_binary(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
1329 respx_mock.post("/foo").mock(
1330 return_value=httpx.Response(200, headers={"Content-Type": "application/binary"}, content='{"foo": "bar"}')
1331 )
1332
1333 response = await async_client.post("/foo", cast_to=httpx.Response)
1334 assert response.status_code == 200
1335 assert isinstance(response, httpx.Response)
1336 assert response.json() == {"foo": "bar"}
1337
1338 def test_copy(self, async_client: AsyncOpenAI) -> None:
1339 copied = async_client.copy()
1340 assert id(copied) != id(async_client)
1341
1342 copied = async_client.copy(api_key="another My API Key")
1343 assert copied.api_key == "another My API Key"
1344 assert async_client.api_key == "My API Key"
1345
1346 copied = async_client.copy(admin_api_key="another My Admin API Key")
1347 assert copied.admin_api_key == "another My Admin API Key"
1348 assert async_client.admin_api_key == "My Admin API Key"
1349
1350 def test_copy_default_options(self, async_client: AsyncOpenAI) -> None:
1351 # options that have a default are overridden correctly
1352 copied = async_client.copy(max_retries=7)
1353 assert copied.max_retries == 7
1354 assert async_client.max_retries == 2
1355
1356 copied2 = copied.copy(max_retries=6)
1357 assert copied2.max_retries == 6
1358 assert copied.max_retries == 7
1359
1360 # timeout
1361 assert isinstance(async_client.timeout, httpx.Timeout)
1362 copied = async_client.copy(timeout=None)
1363 assert copied.timeout is None
1364 assert isinstance(async_client.timeout, httpx.Timeout)
1365
1366 async def test_copy_default_headers(self) -> None:
1367 client = AsyncOpenAI(
1368 base_url=base_url,
1369 api_key=api_key,
1370 admin_api_key=admin_api_key,
1371 _strict_response_validation=True,
1372 default_headers={"X-Foo": "bar"},
1373 )
1374 assert client.default_headers["X-Foo"] == "bar"
1375
1376 # does not override the already given value when not specified
1377 copied = client.copy()
1378 assert copied.default_headers["X-Foo"] == "bar"
1379
1380 # merges already given headers
1381 copied = client.copy(default_headers={"X-Bar": "stainless"})
1382 assert copied.default_headers["X-Foo"] == "bar"
1383 assert copied.default_headers["X-Bar"] == "stainless"
1384
1385 # uses new values for any already given headers
1386 copied = client.copy(default_headers={"X-Foo": "stainless"})
1387 assert copied.default_headers["X-Foo"] == "stainless"
1388
1389 # set_default_headers
1390
1391 # completely overrides already set values
1392 copied = client.copy(set_default_headers={})
1393 assert copied.default_headers.get("X-Foo") is None
1394
1395 copied = client.copy(set_default_headers={"X-Bar": "Robert"})
1396 assert copied.default_headers["X-Bar"] == "Robert"
1397
1398 with pytest.raises(
1399 ValueError,
1400 match="`default_headers` and `set_default_headers` arguments are mutually exclusive",
1401 ):
1402 client.copy(set_default_headers={}, default_headers={"X-Foo": "Bar"})
1403 await client.close()
1404
1405 async def test_copy_default_query(self) -> None:
1406 client = AsyncOpenAI(
1407 base_url=base_url,
1408 api_key=api_key,
1409 admin_api_key=admin_api_key,
1410 _strict_response_validation=True,
1411 default_query={"foo": "bar"},
1412 )
1413 assert _get_params(client)["foo"] == "bar"
1414
1415 # does not override the already given value when not specified
1416 copied = client.copy()
1417 assert _get_params(copied)["foo"] == "bar"
1418
1419 # merges already given params
1420 copied = client.copy(default_query={"bar": "stainless"})
1421 params = _get_params(copied)
1422 assert params["foo"] == "bar"
1423 assert params["bar"] == "stainless"
1424
1425 # uses new values for any already given headers
1426 copied = client.copy(default_query={"foo": "stainless"})
1427 assert _get_params(copied)["foo"] == "stainless"
1428
1429 # set_default_query
1430
1431 # completely overrides already set values
1432 copied = client.copy(set_default_query={})
1433 assert _get_params(copied) == {}
1434
1435 copied = client.copy(set_default_query={"bar": "Robert"})
1436 assert _get_params(copied)["bar"] == "Robert"
1437
1438 with pytest.raises(
1439 ValueError,
1440 # TODO: update
1441 match="`default_query` and `set_default_query` arguments are mutually exclusive",
1442 ):
1443 client.copy(set_default_query={}, default_query={"foo": "Bar"})
1444
1445 await client.close()
1446
1447 def test_copy_signature(self, async_client: AsyncOpenAI) -> None:
1448 # ensure the same parameters that can be passed to the client are defined in the `.copy()` method
1449 init_signature = inspect.signature(
1450 # mypy doesn't like that we access the `__init__` property.
1451 async_client.__init__, # type: ignore[misc]
1452 )
1453 copy_signature = inspect.signature(async_client.copy)
1454 exclude_params = {"transport", "proxies", "_strict_response_validation"}
1455
1456 for name in init_signature.parameters.keys():
1457 if name in exclude_params:
1458 continue
1459
1460 copy_param = copy_signature.parameters.get(name)
1461 assert copy_param is not None, f"copy() signature is missing the {name} param"
1462
1463 @pytest.mark.skipif(sys.version_info >= (3, 10), reason="fails because of a memory leak that started from 3.12")
1464 def test_copy_build_request(self, async_client: AsyncOpenAI) -> None:
1465 options = FinalRequestOptions(method="get", url="/foo")
1466
1467 def build_request(options: FinalRequestOptions) -> None:
1468 client_copy = async_client.copy()
1469 client_copy._build_request(options)
1470
1471 # ensure that the machinery is warmed up before tracing starts.
1472 build_request(options)
1473 gc.collect()
1474
1475 tracemalloc.start(1000)
1476
1477 snapshot_before = tracemalloc.take_snapshot()
1478
1479 ITERATIONS = 10
1480 for _ in range(ITERATIONS):
1481 build_request(options)
1482
1483 gc.collect()
1484 snapshot_after = tracemalloc.take_snapshot()
1485
1486 tracemalloc.stop()
1487
1488 def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
1489 if diff.count == 0:
1490 # Avoid false positives by considering only leaks (i.e. allocations that persist).
1491 return
1492
1493 if diff.count % ITERATIONS != 0:
1494 # Avoid false positives by considering only leaks that appear per iteration.
1495 return
1496
1497 for frame in diff.traceback:
1498 if any(
1499 frame.filename.endswith(fragment)
1500 for fragment in [
1501 # to_raw_response_wrapper leaks through the @functools.wraps() decorator.
1502 #
1503 # removing the decorator fixes the leak for reasons we don't understand.
1504 "openai/_legacy_response.py",
1505 "openai/_response.py",
1506 # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
1507 "openai/_compat.py",
1508 # Standard library leaks we don't care about.
1509 "/logging/__init__.py",
1510 ]
1511 ):
1512 return
1513
1514 leaks.append(diff)
1515
1516 leaks: list[tracemalloc.StatisticDiff] = []
1517 for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
1518 add_leak(leaks, diff)
1519 if leaks:
1520 for leak in leaks:
1521 print("MEMORY LEAK:", leak)
1522 for frame in leak.traceback:
1523 print(frame)
1524 raise AssertionError()
1525
1526 async def test_request_timeout(self, async_client: AsyncOpenAI) -> None:
1527 request = async_client._build_request(FinalRequestOptions(method="get", url="/foo"))
1528 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1529 assert timeout == DEFAULT_TIMEOUT
1530
1531 request = async_client._build_request(
1532 FinalRequestOptions(method="get", url="/foo", timeout=httpx.Timeout(100.0))
1533 )
1534 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1535 assert timeout == httpx.Timeout(100.0)
1536
1537 async def test_client_timeout_option(self) -> None:
1538 client = AsyncOpenAI(
1539 base_url=base_url,
1540 api_key=api_key,
1541 admin_api_key=admin_api_key,
1542 _strict_response_validation=True,
1543 timeout=httpx.Timeout(0),
1544 )
1545
1546 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1547 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1548 assert timeout == httpx.Timeout(0)
1549
1550 await client.close()
1551
1552 async def test_http_client_timeout_option(self) -> None:
1553 # custom timeout given to the httpx client should be used
1554 async with httpx.AsyncClient(timeout=None) as http_client:
1555 client = AsyncOpenAI(
1556 base_url=base_url,
1557 api_key=api_key,
1558 admin_api_key=admin_api_key,
1559 _strict_response_validation=True,
1560 http_client=http_client,
1561 )
1562
1563 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1564 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1565 assert timeout == httpx.Timeout(None)
1566
1567 await client.close()
1568
1569 # no timeout given to the httpx client should not use the httpx default
1570 async with httpx.AsyncClient() as http_client:
1571 client = AsyncOpenAI(
1572 base_url=base_url,
1573 api_key=api_key,
1574 admin_api_key=admin_api_key,
1575 _strict_response_validation=True,
1576 http_client=http_client,
1577 )
1578
1579 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1580 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1581 assert timeout == DEFAULT_TIMEOUT
1582
1583 await client.close()
1584
1585 # explicitly passing the default timeout currently results in it being ignored
1586 async with httpx.AsyncClient(timeout=HTTPX_DEFAULT_TIMEOUT) as http_client:
1587 client = AsyncOpenAI(
1588 base_url=base_url,
1589 api_key=api_key,
1590 admin_api_key=admin_api_key,
1591 _strict_response_validation=True,
1592 http_client=http_client,
1593 )
1594
1595 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1596 timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
1597 assert timeout == DEFAULT_TIMEOUT # our default
1598
1599 await client.close()
1600
1601 def test_invalid_http_client(self) -> None:
1602 with pytest.raises(TypeError, match="Invalid `http_client` arg"):
1603 with httpx.Client() as http_client:
1604 AsyncOpenAI(
1605 base_url=base_url,
1606 api_key=api_key,
1607 admin_api_key=admin_api_key,
1608 _strict_response_validation=True,
1609 http_client=cast(Any, http_client),
1610 )
1611
1612 async def test_default_headers_option(self) -> None:
1613 test_client = AsyncOpenAI(
1614 base_url=base_url,
1615 api_key=api_key,
1616 admin_api_key=admin_api_key,
1617 _strict_response_validation=True,
1618 default_headers={"X-Foo": "bar"},
1619 )
1620 request = test_client._build_request(FinalRequestOptions(method="get", url="/foo"))
1621 assert request.headers.get("x-foo") == "bar"
1622 assert request.headers.get("x-stainless-lang") == "python"
1623
1624 test_client2 = AsyncOpenAI(
1625 base_url=base_url,
1626 api_key=api_key,
1627 admin_api_key=admin_api_key,
1628 _strict_response_validation=True,
1629 default_headers={
1630 "X-Foo": "stainless",
1631 "X-Stainless-Lang": "my-overriding-header",
1632 },
1633 )
1634 request = test_client2._build_request(FinalRequestOptions(method="get", url="/foo"))
1635 assert request.headers.get("x-foo") == "stainless"
1636 assert request.headers.get("x-stainless-lang") == "my-overriding-header"
1637
1638 await test_client.close()
1639 await test_client2.close()
1640
1641 async def test_validate_headers(self) -> None:
1642 client = AsyncOpenAI(
1643 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
1644 )
1645 options = await client._prepare_options(FinalRequestOptions(method="get", url="/foo"))
1646 request = client._build_request(options)
1647 assert request.headers.get("Authorization") == f"Bearer {api_key}"
1648
1649 admin_request = client._build_request(
1650 FinalRequestOptions(
1651 method="get",
1652 url="/organization/projects",
1653 security={"admin_api_key_auth": True},
1654 )
1655 )
1656 assert admin_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
1657
1658 with update_env(**{"OPENAI_API_KEY": Omit()}):
1659 admin_only = AsyncOpenAI(
1660 base_url=base_url,
1661 api_key=None,
1662 admin_api_key=admin_api_key,
1663 _strict_response_validation=True,
1664 )
1665 admin_only_request = admin_only._build_request(
1666 FinalRequestOptions(
1667 method="get",
1668 url="/organization/projects",
1669 security={"admin_api_key_auth": True},
1670 )
1671 )
1672 assert admin_only_request.headers.get("Authorization") == f"Bearer {admin_api_key}"
1673
1674 with pytest.raises(
1675 TypeError,
1676 match="Could not resolve authentication method",
1677 ):
1678 admin_only._build_request(
1679 FinalRequestOptions(
1680 method="post",
1681 url="/responses",
1682 security={"bearer_auth": True},
1683 )
1684 )
1685
1686 with update_env(
1687 **{
1688 "OPENAI_API_KEY": Omit(),
1689 "OPENAI_ADMIN_KEY": Omit(),
1690 }
1691 ):
1692 with pytest.raises(OpenAIError, match="Missing credentials"):
1693 AsyncOpenAI(base_url=base_url, api_key=None, admin_api_key=None, _strict_response_validation=True)
1694
1695 async def test_default_query_option(self) -> None:
1696 client = AsyncOpenAI(
1697 base_url=base_url,
1698 api_key=api_key,
1699 admin_api_key=admin_api_key,
1700 _strict_response_validation=True,
1701 default_query={"query_param": "bar"},
1702 )
1703 request = client._build_request(FinalRequestOptions(method="get", url="/foo"))
1704 url = httpx.URL(request.url)
1705 assert dict(url.params) == {"query_param": "bar"}
1706
1707 request = client._build_request(
1708 FinalRequestOptions(
1709 method="get",
1710 url="/foo",
1711 params={"foo": "baz", "query_param": "overridden"},
1712 )
1713 )
1714 url = httpx.URL(request.url)
1715 assert dict(url.params) == {"foo": "baz", "query_param": "overridden"}
1716
1717 await client.close()
1718
1719 async def test_hardcoded_query_params_in_url(self, async_client: AsyncOpenAI) -> None:
1720 request = async_client._build_request(FinalRequestOptions(method="get", url="/foo?beta=true"))
1721 url = httpx.URL(request.url)
1722 assert dict(url.params) == {"beta": "true"}
1723
1724 request = async_client._build_request(
1725 FinalRequestOptions(
1726 method="get",
1727 url="/foo?beta=true",
1728 params={"limit": "10", "page": "abc"},
1729 )
1730 )
1731 url = httpx.URL(request.url)
1732 assert dict(url.params) == {"beta": "true", "limit": "10", "page": "abc"}
1733
1734 request = async_client._build_request(
1735 FinalRequestOptions(
1736 method="get",
1737 url="/files/a%2Fb?beta=true",
1738 params={"limit": "10"},
1739 )
1740 )
1741 assert request.url.raw_path == b"/files/a%2Fb?beta=true&limit=10"
1742
1743 def test_request_extra_json(self, client: OpenAI) -> None:
1744 request = client._build_request(
1745 FinalRequestOptions(
1746 method="post",
1747 url="/foo",
1748 json_data={"foo": "bar"},
1749 extra_json={"baz": False},
1750 ),
1751 )
1752 data = json.loads(request.content.decode("utf-8"))
1753 assert data == {"foo": "bar", "baz": False}
1754
1755 request = client._build_request(
1756 FinalRequestOptions(
1757 method="post",
1758 url="/foo",
1759 extra_json={"baz": False},
1760 ),
1761 )
1762 data = json.loads(request.content.decode("utf-8"))
1763 assert data == {"baz": False}
1764
1765 # `extra_json` takes priority over `json_data` when keys clash
1766 request = client._build_request(
1767 FinalRequestOptions(
1768 method="post",
1769 url="/foo",
1770 json_data={"foo": "bar", "baz": True},
1771 extra_json={"baz": None},
1772 ),
1773 )
1774 data = json.loads(request.content.decode("utf-8"))
1775 assert data == {"foo": "bar", "baz": None}
1776
1777 def test_request_extra_headers(self, client: OpenAI) -> None:
1778 request = client._build_request(
1779 FinalRequestOptions(
1780 method="post",
1781 url="/foo",
1782 **make_request_options(extra_headers={"X-Foo": "Foo"}),
1783 ),
1784 )
1785 assert request.headers.get("X-Foo") == "Foo"
1786
1787 # `extra_headers` takes priority over `default_headers` when keys clash
1788 request = client.with_options(default_headers={"X-Bar": "true"})._build_request(
1789 FinalRequestOptions(
1790 method="post",
1791 url="/foo",
1792 **make_request_options(
1793 extra_headers={"X-Bar": "false"},
1794 ),
1795 ),
1796 )
1797 assert request.headers.get("X-Bar") == "false"
1798
1799 def test_request_extra_query(self, client: OpenAI) -> None:
1800 request = client._build_request(
1801 FinalRequestOptions(
1802 method="post",
1803 url="/foo",
1804 **make_request_options(
1805 extra_query={"my_query_param": "Foo"},
1806 ),
1807 ),
1808 )
1809 params = dict(request.url.params)
1810 assert params == {"my_query_param": "Foo"}
1811
1812 # if both `query` and `extra_query` are given, they are merged
1813 request = client._build_request(
1814 FinalRequestOptions(
1815 method="post",
1816 url="/foo",
1817 **make_request_options(
1818 query={"bar": "1"},
1819 extra_query={"foo": "2"},
1820 ),
1821 ),
1822 )
1823 params = dict(request.url.params)
1824 assert params == {"bar": "1", "foo": "2"}
1825
1826 # `extra_query` takes priority over `query` when keys clash
1827 request = client._build_request(
1828 FinalRequestOptions(
1829 method="post",
1830 url="/foo",
1831 **make_request_options(
1832 query={"foo": "1"},
1833 extra_query={"foo": "2"},
1834 ),
1835 ),
1836 )
1837 params = dict(request.url.params)
1838 assert params == {"foo": "2"}
1839
1840 def test_multipart_repeating_array(self, async_client: AsyncOpenAI) -> None:
1841 request = async_client._build_request(
1842 FinalRequestOptions.construct(
1843 method="post",
1844 url="/foo",
1845 headers={"Content-Type": "multipart/form-data; boundary=6b7ba517decee4a450543ea6ae821c82"},
1846 json_data={"array": ["foo", "bar"]},
1847 files=[("foo.txt", b"hello world")],
1848 )
1849 )
1850
1851 assert request.read().split(b"\r\n") == [
1852 b"--6b7ba517decee4a450543ea6ae821c82",
1853 b'Content-Disposition: form-data; name="array[]"',
1854 b"",
1855 b"foo",
1856 b"--6b7ba517decee4a450543ea6ae821c82",
1857 b'Content-Disposition: form-data; name="array[]"',
1858 b"",
1859 b"bar",
1860 b"--6b7ba517decee4a450543ea6ae821c82",
1861 b'Content-Disposition: form-data; name="foo.txt"; filename="upload"',
1862 b"Content-Type: application/octet-stream",
1863 b"",
1864 b"hello world",
1865 b"--6b7ba517decee4a450543ea6ae821c82--",
1866 b"",
1867 ]
1868
1869 @pytest.mark.respx(base_url=base_url)
1870 async def test_binary_content_upload(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
1871 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
1872
1873 file_content = b"Hello, this is a test file."
1874
1875 response = await async_client.post(
1876 "/upload",
1877 content=file_content,
1878 cast_to=httpx.Response,
1879 options={"headers": {"Content-Type": "application/octet-stream"}},
1880 )
1881
1882 assert response.status_code == 200
1883 assert response.request.headers["Content-Type"] == "application/octet-stream"
1884 assert response.content == file_content
1885
1886 async def test_binary_content_upload_with_asynciterator(self) -> None:
1887 file_content = b"Hello, this is a test file."
1888 counter = Counter()
1889 iterator = _make_async_iterator([file_content], counter=counter)
1890
1891 async def mock_handler(request: httpx.Request) -> httpx.Response:
1892 assert counter.value == 0, "the request body should not have been read"
1893 return httpx.Response(200, content=await request.aread())
1894
1895 async with AsyncOpenAI(
1896 base_url=base_url,
1897 api_key=api_key,
1898 admin_api_key=admin_api_key,
1899 _strict_response_validation=True,
1900 http_client=httpx.AsyncClient(transport=MockTransport(handler=mock_handler)),
1901 ) as client:
1902 response = await client.post(
1903 "/upload",
1904 content=iterator,
1905 cast_to=httpx.Response,
1906 options={"headers": {"Content-Type": "application/octet-stream"}},
1907 )
1908
1909 assert response.status_code == 200
1910 assert response.request.headers["Content-Type"] == "application/octet-stream"
1911 assert response.content == file_content
1912 assert counter.value == 1
1913
1914 @pytest.mark.respx(base_url=base_url)
1915 async def test_binary_content_upload_with_body_is_deprecated(
1916 self, respx_mock: MockRouter, async_client: AsyncOpenAI
1917 ) -> None:
1918 respx_mock.post("/upload").mock(side_effect=mirror_request_content)
1919
1920 file_content = b"Hello, this is a test file."
1921
1922 with pytest.deprecated_call(
1923 match="Passing raw bytes as `body` is deprecated and will be removed in a future version. Please pass raw bytes via the `content` parameter instead."
1924 ):
1925 response = await async_client.post(
1926 "/upload",
1927 body=file_content,
1928 cast_to=httpx.Response,
1929 options={"headers": {"Content-Type": "application/octet-stream"}},
1930 )
1931
1932 assert response.status_code == 200
1933 assert response.request.headers["Content-Type"] == "application/octet-stream"
1934 assert response.content == file_content
1935
1936 @pytest.mark.respx(base_url=base_url)
1937 async def test_basic_union_response(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
1938 class Model1(BaseModel):
1939 name: str
1940
1941 class Model2(BaseModel):
1942 foo: str
1943
1944 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
1945
1946 response = await async_client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
1947 assert isinstance(response, Model2)
1948 assert response.foo == "bar"
1949
1950 @pytest.mark.respx(base_url=base_url)
1951 async def test_union_response_different_types(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
1952 """Union of objects with the same field name using a different type"""
1953
1954 class Model1(BaseModel):
1955 foo: int
1956
1957 class Model2(BaseModel):
1958 foo: str
1959
1960 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
1961
1962 response = await async_client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
1963 assert isinstance(response, Model2)
1964 assert response.foo == "bar"
1965
1966 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": 1}))
1967
1968 response = await async_client.get("/foo", cast_to=cast(Any, Union[Model1, Model2]))
1969 assert isinstance(response, Model1)
1970 assert response.foo == 1
1971
1972 @pytest.mark.respx(base_url=base_url)
1973 async def test_non_application_json_content_type_for_json_data(
1974 self, respx_mock: MockRouter, async_client: AsyncOpenAI
1975 ) -> None:
1976 """
1977 Response that sets Content-Type to something other than application/json but returns json data
1978 """
1979
1980 class Model(BaseModel):
1981 foo: int
1982
1983 respx_mock.get("/foo").mock(
1984 return_value=httpx.Response(
1985 200,
1986 content=json.dumps({"foo": 2}),
1987 headers={"Content-Type": "application/text"},
1988 )
1989 )
1990
1991 response = await async_client.get("/foo", cast_to=Model)
1992 assert isinstance(response, Model)
1993 assert response.foo == 2
1994
1995 async def test_base_url_setter(self) -> None:
1996 client = AsyncOpenAI(
1997 base_url="https://example.com/from_init",
1998 api_key=api_key,
1999 admin_api_key=admin_api_key,
2000 _strict_response_validation=True,
2001 )
2002 assert client.base_url == "https://example.com/from_init/"
2003
2004 client.base_url = "https://example.com/from_setter" # type: ignore[assignment]
2005
2006 assert client.base_url == "https://example.com/from_setter/"
2007
2008 await client.close()
2009
2010 async def test_base_url_env(self) -> None:
2011 with update_env(OPENAI_BASE_URL="http://localhost:5000/from/env"):
2012 client = AsyncOpenAI(api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True)
2013 assert client.base_url == "http://localhost:5000/from/env/"
2014
2015 @pytest.mark.parametrize(
2016 "client",
2017 [
2018 AsyncOpenAI(
2019 base_url="http://localhost:5000/custom/path/",
2020 api_key=api_key,
2021 admin_api_key=admin_api_key,
2022 _strict_response_validation=True,
2023 ),
2024 AsyncOpenAI(
2025 base_url="http://localhost:5000/custom/path/",
2026 api_key=api_key,
2027 admin_api_key=admin_api_key,
2028 _strict_response_validation=True,
2029 http_client=httpx.AsyncClient(),
2030 ),
2031 ],
2032 ids=["standard", "custom http client"],
2033 )
2034 async def test_base_url_trailing_slash(self, client: AsyncOpenAI) -> None:
2035 request = client._build_request(
2036 FinalRequestOptions(
2037 method="post",
2038 url="/foo",
2039 json_data={"foo": "bar"},
2040 ),
2041 )
2042 assert request.url == "http://localhost:5000/custom/path/foo"
2043 await client.close()
2044
2045 @pytest.mark.parametrize(
2046 "client",
2047 [
2048 AsyncOpenAI(
2049 base_url="http://localhost:5000/custom/path/",
2050 api_key=api_key,
2051 admin_api_key=admin_api_key,
2052 _strict_response_validation=True,
2053 ),
2054 AsyncOpenAI(
2055 base_url="http://localhost:5000/custom/path/",
2056 api_key=api_key,
2057 admin_api_key=admin_api_key,
2058 _strict_response_validation=True,
2059 http_client=httpx.AsyncClient(),
2060 ),
2061 ],
2062 ids=["standard", "custom http client"],
2063 )
2064 async def test_base_url_no_trailing_slash(self, client: AsyncOpenAI) -> None:
2065 request = client._build_request(
2066 FinalRequestOptions(
2067 method="post",
2068 url="/foo",
2069 json_data={"foo": "bar"},
2070 ),
2071 )
2072 assert request.url == "http://localhost:5000/custom/path/foo"
2073 await client.close()
2074
2075 @pytest.mark.parametrize(
2076 "client",
2077 [
2078 AsyncOpenAI(
2079 base_url="http://localhost:5000/custom/path/",
2080 api_key=api_key,
2081 admin_api_key=admin_api_key,
2082 _strict_response_validation=True,
2083 ),
2084 AsyncOpenAI(
2085 base_url="http://localhost:5000/custom/path/",
2086 api_key=api_key,
2087 admin_api_key=admin_api_key,
2088 _strict_response_validation=True,
2089 http_client=httpx.AsyncClient(),
2090 ),
2091 ],
2092 ids=["standard", "custom http client"],
2093 )
2094 async def test_absolute_request_url(self, client: AsyncOpenAI) -> None:
2095 request = client._build_request(
2096 FinalRequestOptions(
2097 method="post",
2098 url="https://myapi.com/foo",
2099 json_data={"foo": "bar"},
2100 ),
2101 )
2102 assert request.url == "https://myapi.com/foo"
2103 await client.close()
2104
2105 async def test_copied_client_does_not_close_http(self) -> None:
2106 test_client = AsyncOpenAI(
2107 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
2108 )
2109 assert not test_client.is_closed()
2110
2111 copied = test_client.copy()
2112 assert copied is not test_client
2113
2114 del copied
2115
2116 await asyncio.sleep(0.2)
2117 assert not test_client.is_closed()
2118
2119 async def test_client_context_manager(self) -> None:
2120 test_client = AsyncOpenAI(
2121 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
2122 )
2123 async with test_client as c2:
2124 assert c2 is test_client
2125 assert not c2.is_closed()
2126 assert not test_client.is_closed()
2127 assert test_client.is_closed()
2128
2129 @pytest.mark.respx(base_url=base_url)
2130 async def test_client_response_validation_error(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2131 class Model(BaseModel):
2132 foo: str
2133
2134 respx_mock.get("/foo").mock(return_value=httpx.Response(200, json={"foo": {"invalid": True}}))
2135
2136 with pytest.raises(APIResponseValidationError) as exc:
2137 await async_client.get("/foo", cast_to=Model)
2138
2139 assert isinstance(exc.value.__cause__, ValidationError)
2140
2141 async def test_client_max_retries_validation(self) -> None:
2142 with pytest.raises(TypeError, match=r"max_retries cannot be None"):
2143 AsyncOpenAI(
2144 base_url=base_url,
2145 api_key=api_key,
2146 admin_api_key=admin_api_key,
2147 _strict_response_validation=True,
2148 max_retries=cast(Any, None),
2149 )
2150
2151 @pytest.mark.respx(base_url=base_url)
2152 async def test_default_stream_cls(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2153 class Model(BaseModel):
2154 name: str
2155
2156 respx_mock.post("/foo").mock(return_value=httpx.Response(200, json={"foo": "bar"}))
2157
2158 stream = await async_client.post("/foo", cast_to=Model, stream=True, stream_cls=AsyncStream[Model])
2159 assert isinstance(stream, AsyncStream)
2160 await stream.response.aclose()
2161
2162 @pytest.mark.respx(base_url=base_url)
2163 async def test_received_text_for_expected_json(self, respx_mock: MockRouter) -> None:
2164 class Model(BaseModel):
2165 name: str
2166
2167 respx_mock.get("/foo").mock(return_value=httpx.Response(200, text="my-custom-format"))
2168
2169 strict_client = AsyncOpenAI(
2170 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=True
2171 )
2172
2173 with pytest.raises(APIResponseValidationError):
2174 await strict_client.get("/foo", cast_to=Model)
2175
2176 non_strict_client = AsyncOpenAI(
2177 base_url=base_url, api_key=api_key, admin_api_key=admin_api_key, _strict_response_validation=False
2178 )
2179
2180 response = await non_strict_client.get("/foo", cast_to=Model)
2181 assert isinstance(response, str) # type: ignore[unreachable]
2182
2183 await strict_client.close()
2184 await non_strict_client.close()
2185
2186 @pytest.mark.parametrize(
2187 "remaining_retries,retry_after,timeout",
2188 [
2189 [3, "20", 20],
2190 [3, "0", 0.5],
2191 [3, "-10", 0.5],
2192 [3, "60", 60],
2193 [3, "61", 0.5],
2194 [3, "Fri, 29 Sep 2023 16:26:57 GMT", 20],
2195 [3, "Fri, 29 Sep 2023 16:26:37 GMT", 0.5],
2196 [3, "Fri, 29 Sep 2023 16:26:27 GMT", 0.5],
2197 [3, "Fri, 29 Sep 2023 16:27:37 GMT", 60],
2198 [3, "Fri, 29 Sep 2023 16:27:38 GMT", 0.5],
2199 [3, "99999999999999999999999999999999999", 0.5],
2200 [3, "Zun, 29 Sep 2023 16:26:27 GMT", 0.5],
2201 [3, "", 0.5],
2202 [2, "", 0.5 * 2.0],
2203 [1, "", 0.5 * 4.0],
2204 [-1100, "", 8], # test large number potentially overflowing
2205 ],
2206 )
2207 @mock.patch("time.time", mock.MagicMock(return_value=1696004797))
2208 async def test_parse_retry_after_header(
2209 self, remaining_retries: int, retry_after: str, timeout: float, async_client: AsyncOpenAI
2210 ) -> None:
2211 headers = httpx.Headers({"retry-after": retry_after})
2212 options = FinalRequestOptions(method="get", url="/foo", max_retries=3)
2213 calculated = async_client._calculate_retry_timeout(remaining_retries, options, headers)
2214 assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType]
2215
2216 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2217 @pytest.mark.respx(base_url=base_url)
2218 async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2219 respx_mock.post("/chat/completions").mock(side_effect=httpx.TimeoutException("Test timeout error"))
2220
2221 with pytest.raises(APITimeoutError):
2222 await async_client.chat.completions.with_streaming_response.create(
2223 messages=[
2224 {
2225 "content": "string",
2226 "role": "developer",
2227 }
2228 ],
2229 model="gpt-5.4",
2230 ).__aenter__()
2231
2232 assert _get_open_connections(async_client) == 0
2233
2234 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2235 @pytest.mark.respx(base_url=base_url)
2236 async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2237 respx_mock.post("/chat/completions").mock(return_value=httpx.Response(500))
2238
2239 with pytest.raises(APIStatusError):
2240 await async_client.chat.completions.with_streaming_response.create(
2241 messages=[
2242 {
2243 "content": "string",
2244 "role": "developer",
2245 }
2246 ],
2247 model="gpt-5.4",
2248 ).__aenter__()
2249 assert _get_open_connections(async_client) == 0
2250
2251 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2252 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2253 @pytest.mark.respx(base_url=base_url)
2254 @pytest.mark.parametrize("failure_mode", ["status", "exception"])
2255 async def test_retries_taken(
2256 self,
2257 async_client: AsyncOpenAI,
2258 failures_before_success: int,
2259 failure_mode: Literal["status", "exception"],
2260 respx_mock: MockRouter,
2261 ) -> None:
2262 client = async_client.with_options(max_retries=4)
2263
2264 nb_retries = 0
2265
2266 def retry_handler(_request: httpx.Request) -> httpx.Response:
2267 nonlocal nb_retries
2268 if nb_retries < failures_before_success:
2269 nb_retries += 1
2270 if failure_mode == "exception":
2271 raise RuntimeError("oops")
2272 return httpx.Response(500)
2273 return httpx.Response(200)
2274
2275 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2276
2277 response = await client.chat.completions.with_raw_response.create(
2278 messages=[
2279 {
2280 "content": "string",
2281 "role": "developer",
2282 }
2283 ],
2284 model="gpt-5.4",
2285 )
2286
2287 assert response.retries_taken == failures_before_success
2288 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
2289
2290 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2291 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2292 @pytest.mark.respx(base_url=base_url)
2293 async def test_omit_retry_count_header(
2294 self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
2295 ) -> None:
2296 client = async_client.with_options(max_retries=4)
2297
2298 nb_retries = 0
2299
2300 def retry_handler(_request: httpx.Request) -> httpx.Response:
2301 nonlocal nb_retries
2302 if nb_retries < failures_before_success:
2303 nb_retries += 1
2304 return httpx.Response(500)
2305 return httpx.Response(200)
2306
2307 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2308
2309 response = await client.chat.completions.with_raw_response.create(
2310 messages=[
2311 {
2312 "content": "string",
2313 "role": "developer",
2314 }
2315 ],
2316 model="gpt-5.4",
2317 extra_headers={"x-stainless-retry-count": Omit()},
2318 )
2319
2320 assert len(response.http_request.headers.get_list("x-stainless-retry-count")) == 0
2321
2322 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2323 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2324 @pytest.mark.respx(base_url=base_url)
2325 async def test_overwrite_retry_count_header(
2326 self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
2327 ) -> None:
2328 client = async_client.with_options(max_retries=4)
2329
2330 nb_retries = 0
2331
2332 def retry_handler(_request: httpx.Request) -> httpx.Response:
2333 nonlocal nb_retries
2334 if nb_retries < failures_before_success:
2335 nb_retries += 1
2336 return httpx.Response(500)
2337 return httpx.Response(200)
2338
2339 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2340
2341 response = await client.chat.completions.with_raw_response.create(
2342 messages=[
2343 {
2344 "content": "string",
2345 "role": "developer",
2346 }
2347 ],
2348 model="gpt-5.4",
2349 extra_headers={"x-stainless-retry-count": "42"},
2350 )
2351
2352 assert response.http_request.headers.get("x-stainless-retry-count") == "42"
2353
2354 @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
2355 @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
2356 @pytest.mark.respx(base_url=base_url)
2357 async def test_retries_taken_new_response_class(
2358 self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
2359 ) -> None:
2360 client = async_client.with_options(max_retries=4)
2361
2362 nb_retries = 0
2363
2364 def retry_handler(_request: httpx.Request) -> httpx.Response:
2365 nonlocal nb_retries
2366 if nb_retries < failures_before_success:
2367 nb_retries += 1
2368 return httpx.Response(500)
2369 return httpx.Response(200)
2370
2371 respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
2372
2373 async with client.chat.completions.with_streaming_response.create(
2374 messages=[
2375 {
2376 "content": "string",
2377 "role": "developer",
2378 }
2379 ],
2380 model="gpt-5.4",
2381 ) as response:
2382 assert response.retries_taken == failures_before_success
2383 assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
2384
2385 async def test_get_platform(self) -> None:
2386 platform = await asyncify(get_platform)()
2387 assert isinstance(platform, (str, OtherPlatform))
2388
2389 async def test_proxy_environment_variables(self, monkeypatch: pytest.MonkeyPatch) -> None:
2390 # Test that the proxy environment variables are set correctly
2391 monkeypatch.setenv("HTTPS_PROXY", "https://example.org")
2392 # Delete in case our environment has any proxy env vars set
2393 monkeypatch.delenv("HTTP_PROXY", raising=False)
2394 monkeypatch.delenv("ALL_PROXY", raising=False)
2395 monkeypatch.delenv("NO_PROXY", raising=False)
2396 monkeypatch.delenv("http_proxy", raising=False)
2397 monkeypatch.delenv("https_proxy", raising=False)
2398 monkeypatch.delenv("all_proxy", raising=False)
2399 monkeypatch.delenv("no_proxy", raising=False)
2400
2401 client = DefaultAsyncHttpxClient()
2402
2403 mounts = tuple(client._mounts.items())
2404 assert len(mounts) == 1
2405 assert mounts[0][0].pattern == "https://"
2406
2407 @pytest.mark.filterwarnings("ignore:.*deprecated.*:DeprecationWarning")
2408 async def test_default_client_creation(self) -> None:
2409 # Ensure that the client can be initialized without any exceptions
2410 DefaultAsyncHttpxClient(
2411 verify=True,
2412 cert=None,
2413 trust_env=True,
2414 http1=True,
2415 http2=False,
2416 limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
2417 )
2418
2419 @pytest.mark.respx(base_url=base_url)
2420 async def test_follow_redirects(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2421 # Test that the default follow_redirects=True allows following redirects
2422 respx_mock.post("/redirect").mock(
2423 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
2424 )
2425 respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
2426
2427 response = await async_client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
2428 assert response.status_code == 200
2429 assert response.json() == {"status": "ok"}
2430
2431 @pytest.mark.respx(base_url=base_url)
2432 async def test_follow_redirects_disabled(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
2433 # Test that follow_redirects=False prevents following redirects
2434 respx_mock.post("/redirect").mock(
2435 return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
2436 )
2437
2438 with pytest.raises(APIStatusError) as exc_info:
2439 await async_client.post(
2440 "/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
2441 )
2442
2443 assert exc_info.value.response.status_code == 302
2444 assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
2445
2446 async def test_api_key_before_after_refresh_provider(self) -> None:
2447 async def mock_api_key_provider():
2448 return "test_bearer_token"
2449
2450 client = AsyncOpenAI(base_url=base_url, api_key=mock_api_key_provider)
2451
2452 assert client.api_key == ""
2453 assert "Authorization" not in client.auth_headers
2454
2455 await client._refresh_api_key()
2456
2457 assert client.api_key == "test_bearer_token"
2458 assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
2459
2460 async def test_api_key_before_after_refresh_str(self) -> None:
2461 client = AsyncOpenAI(base_url=base_url, api_key="test_api_key")
2462
2463 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
2464 await client._refresh_api_key()
2465
2466 assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
2467
2468 @pytest.mark.respx()
2469 async def test_bearer_token_refresh_async(self, respx_mock: MockRouter) -> None:
2470 respx_mock.post(base_url + "/chat/completions").mock(
2471 side_effect=[
2472 httpx.Response(500, json={"error": "server error"}),
2473 httpx.Response(200, json={"foo": "bar"}),
2474 ]
2475 )
2476
2477 counter = 0
2478
2479 async def token_provider() -> str:
2480 nonlocal counter
2481
2482 counter += 1
2483
2484 if counter == 1:
2485 return "first"
2486
2487 return "second"
2488
2489 client = AsyncOpenAI(base_url=base_url, api_key=token_provider)
2490 await client.chat.completions.create(messages=[], model="gpt-4")
2491
2492 calls = cast("list[MockRequestCall]", respx_mock.calls)
2493 assert len(calls) == 2
2494
2495 assert calls[0].request.headers.get("Authorization") == "Bearer first"
2496 assert calls[1].request.headers.get("Authorization") == "Bearer second"
2497
2498 async def test_copy_auth(self) -> None:
2499 async def token_provider_1() -> str:
2500 return "test_bearer_token_1"
2501
2502 async def token_provider_2() -> str:
2503 return "test_bearer_token_2"
2504
2505 client = AsyncOpenAI(base_url=base_url, api_key=token_provider_1).copy(api_key=token_provider_2)
2506 await client._refresh_api_key()
2507 assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}
2508
2509
2510class TestWorkloadIdentity401Retry:
2511 @pytest.mark.respx()
2512 def test_workload_identity_401_retry(self, respx_mock: MockRouter) -> None:
2513 provider_call_count = 0
2514
2515 def provider() -> str:
2516 nonlocal provider_call_count
2517 provider_call_count += 1
2518 return f"external-subject-token-{provider_call_count}"
2519
2520 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2521 side_effect=[
2522 httpx.Response(
2523 200,
2524 json={
2525 "access_token": "openai-access-token-1",
2526 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2527 "token_type": "Bearer",
2528 "expires_in": 3600,
2529 },
2530 ),
2531 httpx.Response(
2532 200,
2533 json={
2534 "access_token": "openai-access-token-2",
2535 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2536 "token_type": "Bearer",
2537 "expires_in": 3600,
2538 },
2539 ),
2540 ]
2541 )
2542
2543 respx_mock.post(base_url + "/chat/completions").mock(
2544 side_effect=[
2545 httpx.Response(401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}),
2546 httpx.Response(
2547 200,
2548 json={
2549 "id": "chatcmpl-123",
2550 "object": "chat.completion",
2551 "created": 1234567890,
2552 "model": "gpt-4",
2553 "choices": [],
2554 },
2555 ),
2556 ]
2557 )
2558
2559 with OpenAI(
2560 base_url=base_url,
2561 workload_identity={
2562 **workload_identity,
2563 "provider": {
2564 "get_token": provider,
2565 "token_type": "jwt",
2566 },
2567 },
2568 organization="org_123",
2569 project="proj_123",
2570 _strict_response_validation=True,
2571 ) as client:
2572 client.chat.completions.create(messages=[], model="gpt-4")
2573
2574 calls = cast("list[MockRequestCall]", respx_mock.calls)
2575 assert len(calls) == 4
2576
2577 assert calls[0].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2578 assert calls[1].request.url == httpx.URL(base_url + "/chat/completions")
2579 assert calls[1].request.headers.get("Authorization") == "Bearer openai-access-token-1"
2580
2581 assert calls[2].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2582
2583 assert calls[3].request.url == httpx.URL(base_url + "/chat/completions")
2584 assert calls[3].request.headers.get("Authorization") == "Bearer openai-access-token-2"
2585
2586 assert provider_call_count == 2
2587
2588 @pytest.mark.respx()
2589 def test_401_without_workload_identity_no_retry(self, respx_mock: MockRouter) -> None:
2590 respx_mock.post(base_url + "/chat/completions").mock(
2591 return_value=httpx.Response(
2592 401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}
2593 )
2594 )
2595
2596 with OpenAI(
2597 base_url=base_url,
2598 api_key="test-api-key",
2599 _strict_response_validation=True,
2600 ) as client:
2601 with pytest.raises(APIStatusError) as exc_info:
2602 client.chat.completions.create(messages=[], model="gpt-4")
2603
2604 assert exc_info.value.status_code == 401
2605
2606 calls = cast("list[MockRequestCall]", respx_mock.calls)
2607 assert len(calls) == 1
2608
2609 @pytest.mark.respx()
2610 def test_non_401_errors_no_retry(self, respx_mock: MockRouter) -> None:
2611 provider_call_count = 0
2612
2613 def provider() -> str:
2614 nonlocal provider_call_count
2615 provider_call_count += 1
2616 return "external-subject-token"
2617
2618 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2619 return_value=httpx.Response(
2620 200,
2621 json={
2622 "access_token": "openai-access-token-1",
2623 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2624 "token_type": "Bearer",
2625 "expires_in": 3600,
2626 },
2627 )
2628 )
2629
2630 respx_mock.post(base_url + "/chat/completions").mock(
2631 return_value=httpx.Response(403, json={"error": {"message": "Forbidden", "type": "invalid_request_error"}})
2632 )
2633
2634 with OpenAI(
2635 base_url=base_url,
2636 workload_identity={
2637 **workload_identity,
2638 "provider": {
2639 "get_token": provider,
2640 "token_type": "jwt",
2641 },
2642 },
2643 organization="org_123",
2644 project="proj_123",
2645 _strict_response_validation=True,
2646 ) as client:
2647 with pytest.raises(APIStatusError) as exc_info:
2648 client.chat.completions.create(messages=[], model="gpt-4")
2649
2650 assert exc_info.value.status_code == 403
2651
2652 calls = cast("list[MockRequestCall]", respx_mock.calls)
2653 assert len(calls) == 2
2654
2655 assert provider_call_count == 1
2656
2657
2658class TestAsyncWorkloadIdentity401Retry:
2659 @pytest.mark.respx()
2660 async def test_workload_identity_401_retry(self, respx_mock: MockRouter) -> None:
2661 provider_call_count = 0
2662
2663 def provider() -> str:
2664 nonlocal provider_call_count
2665 provider_call_count += 1
2666 return f"external-subject-token-{provider_call_count}"
2667
2668 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2669 side_effect=[
2670 httpx.Response(
2671 200,
2672 json={
2673 "access_token": "openai-access-token-1",
2674 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2675 "token_type": "Bearer",
2676 "expires_in": 3600,
2677 },
2678 ),
2679 httpx.Response(
2680 200,
2681 json={
2682 "access_token": "openai-access-token-2",
2683 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2684 "token_type": "Bearer",
2685 "expires_in": 3600,
2686 },
2687 ),
2688 ]
2689 )
2690
2691 respx_mock.post(base_url + "/chat/completions").mock(
2692 side_effect=[
2693 httpx.Response(401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}),
2694 httpx.Response(
2695 200,
2696 json={
2697 "id": "chatcmpl-123",
2698 "object": "chat.completion",
2699 "created": 1234567890,
2700 "model": "gpt-4",
2701 "choices": [],
2702 },
2703 ),
2704 ]
2705 )
2706
2707 async with AsyncOpenAI(
2708 base_url=base_url,
2709 workload_identity={
2710 **workload_identity,
2711 "provider": {
2712 "get_token": provider,
2713 "token_type": "jwt",
2714 },
2715 },
2716 organization="org_123",
2717 project="proj_123",
2718 _strict_response_validation=True,
2719 ) as client:
2720 await client.chat.completions.create(messages=[], model="gpt-4")
2721
2722 calls = cast("list[MockRequestCall]", respx_mock.calls)
2723 assert len(calls) == 4
2724
2725 assert calls[0].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2726 assert calls[1].request.url == httpx.URL(base_url + "/chat/completions")
2727 assert calls[1].request.headers.get("Authorization") == "Bearer openai-access-token-1"
2728
2729 assert calls[2].request.url == httpx.URL("https://auth.openai.com/oauth/token")
2730
2731 assert calls[3].request.url == httpx.URL(base_url + "/chat/completions")
2732 assert calls[3].request.headers.get("Authorization") == "Bearer openai-access-token-2"
2733
2734 assert provider_call_count == 2
2735
2736 @pytest.mark.respx()
2737 async def test_401_without_workload_identity_no_retry(self, respx_mock: MockRouter) -> None:
2738 respx_mock.post(base_url + "/chat/completions").mock(
2739 return_value=httpx.Response(
2740 401, json={"error": {"message": "Unauthorized", "type": "invalid_request_error"}}
2741 )
2742 )
2743
2744 async with AsyncOpenAI(
2745 base_url=base_url,
2746 api_key="test-api-key",
2747 _strict_response_validation=True,
2748 ) as client:
2749 with pytest.raises(APIStatusError) as exc_info:
2750 await client.chat.completions.create(messages=[], model="gpt-4")
2751
2752 assert exc_info.value.status_code == 401
2753
2754 calls = cast("list[MockRequestCall]", respx_mock.calls)
2755 assert len(calls) == 1
2756
2757 @pytest.mark.respx()
2758 async def test_non_401_errors_no_retry(self, respx_mock: MockRouter) -> None:
2759 provider_call_count = 0
2760
2761 def provider() -> str:
2762 nonlocal provider_call_count
2763 provider_call_count += 1
2764 return "external-subject-token"
2765
2766 respx_mock.post("https://auth.openai.com/oauth/token").mock(
2767 return_value=httpx.Response(
2768 200,
2769 json={
2770 "access_token": "openai-access-token-1",
2771 "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
2772 "token_type": "Bearer",
2773 "expires_in": 3600,
2774 },
2775 )
2776 )
2777
2778 respx_mock.post(base_url + "/chat/completions").mock(
2779 return_value=httpx.Response(403, json={"error": {"message": "Forbidden", "type": "invalid_request_error"}})
2780 )
2781
2782 async with AsyncOpenAI(
2783 base_url=base_url,
2784 workload_identity={
2785 **workload_identity,
2786 "provider": {
2787 "get_token": provider,
2788 "token_type": "jwt",
2789 },
2790 },
2791 organization="org_123",
2792 project="proj_123",
2793 _strict_response_validation=True,
2794 ) as client:
2795 with pytest.raises(APIStatusError) as exc_info:
2796 await client.chat.completions.create(messages=[], model="gpt-4")
2797
2798 assert exc_info.value.status_code == 403
2799
2800 calls = cast("list[MockRequestCall]", respx_mock.calls)
2801 assert len(calls) == 2
2802
2803 assert provider_call_count == 1