prefect-client 2.14.21__py3-none-any.whl → 2.15.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (103) hide show
  1. prefect/_internal/concurrency/api.py +37 -2
  2. prefect/_internal/concurrency/calls.py +9 -0
  3. prefect/_internal/concurrency/cancellation.py +3 -1
  4. prefect/_internal/concurrency/event_loop.py +2 -2
  5. prefect/_internal/concurrency/threads.py +3 -2
  6. prefect/_internal/pydantic/annotations/pendulum.py +4 -4
  7. prefect/_internal/pydantic/v2_schema.py +2 -2
  8. prefect/_vendor/fastapi/__init__.py +1 -1
  9. prefect/_vendor/fastapi/applications.py +13 -13
  10. prefect/_vendor/fastapi/background.py +3 -1
  11. prefect/_vendor/fastapi/concurrency.py +7 -3
  12. prefect/_vendor/fastapi/datastructures.py +9 -7
  13. prefect/_vendor/fastapi/dependencies/utils.py +12 -7
  14. prefect/_vendor/fastapi/encoders.py +1 -1
  15. prefect/_vendor/fastapi/exception_handlers.py +7 -4
  16. prefect/_vendor/fastapi/exceptions.py +4 -2
  17. prefect/_vendor/fastapi/middleware/__init__.py +1 -1
  18. prefect/_vendor/fastapi/middleware/asyncexitstack.py +1 -1
  19. prefect/_vendor/fastapi/middleware/cors.py +3 -1
  20. prefect/_vendor/fastapi/middleware/gzip.py +3 -1
  21. prefect/_vendor/fastapi/middleware/httpsredirect.py +1 -1
  22. prefect/_vendor/fastapi/middleware/trustedhost.py +1 -1
  23. prefect/_vendor/fastapi/middleware/wsgi.py +3 -1
  24. prefect/_vendor/fastapi/openapi/docs.py +1 -1
  25. prefect/_vendor/fastapi/openapi/utils.py +3 -3
  26. prefect/_vendor/fastapi/requests.py +4 -2
  27. prefect/_vendor/fastapi/responses.py +13 -7
  28. prefect/_vendor/fastapi/routing.py +15 -15
  29. prefect/_vendor/fastapi/security/api_key.py +3 -3
  30. prefect/_vendor/fastapi/security/http.py +2 -2
  31. prefect/_vendor/fastapi/security/oauth2.py +2 -2
  32. prefect/_vendor/fastapi/security/open_id_connect_url.py +3 -3
  33. prefect/_vendor/fastapi/staticfiles.py +1 -1
  34. prefect/_vendor/fastapi/templating.py +3 -1
  35. prefect/_vendor/fastapi/testclient.py +1 -1
  36. prefect/_vendor/fastapi/utils.py +3 -3
  37. prefect/_vendor/fastapi/websockets.py +7 -3
  38. prefect/_vendor/starlette/__init__.py +1 -0
  39. prefect/_vendor/starlette/_compat.py +28 -0
  40. prefect/_vendor/starlette/_exception_handler.py +80 -0
  41. prefect/_vendor/starlette/_utils.py +88 -0
  42. prefect/_vendor/starlette/applications.py +261 -0
  43. prefect/_vendor/starlette/authentication.py +159 -0
  44. prefect/_vendor/starlette/background.py +43 -0
  45. prefect/_vendor/starlette/concurrency.py +59 -0
  46. prefect/_vendor/starlette/config.py +151 -0
  47. prefect/_vendor/starlette/convertors.py +87 -0
  48. prefect/_vendor/starlette/datastructures.py +707 -0
  49. prefect/_vendor/starlette/endpoints.py +130 -0
  50. prefect/_vendor/starlette/exceptions.py +60 -0
  51. prefect/_vendor/starlette/formparsers.py +276 -0
  52. prefect/_vendor/starlette/middleware/__init__.py +17 -0
  53. prefect/_vendor/starlette/middleware/authentication.py +52 -0
  54. prefect/_vendor/starlette/middleware/base.py +220 -0
  55. prefect/_vendor/starlette/middleware/cors.py +176 -0
  56. prefect/_vendor/starlette/middleware/errors.py +265 -0
  57. prefect/_vendor/starlette/middleware/exceptions.py +74 -0
  58. prefect/_vendor/starlette/middleware/gzip.py +113 -0
  59. prefect/_vendor/starlette/middleware/httpsredirect.py +19 -0
  60. prefect/_vendor/starlette/middleware/sessions.py +82 -0
  61. prefect/_vendor/starlette/middleware/trustedhost.py +64 -0
  62. prefect/_vendor/starlette/middleware/wsgi.py +147 -0
  63. prefect/_vendor/starlette/requests.py +328 -0
  64. prefect/_vendor/starlette/responses.py +347 -0
  65. prefect/_vendor/starlette/routing.py +933 -0
  66. prefect/_vendor/starlette/schemas.py +154 -0
  67. prefect/_vendor/starlette/staticfiles.py +248 -0
  68. prefect/_vendor/starlette/status.py +199 -0
  69. prefect/_vendor/starlette/templating.py +231 -0
  70. prefect/_vendor/starlette/testclient.py +805 -0
  71. prefect/_vendor/starlette/types.py +30 -0
  72. prefect/_vendor/starlette/websockets.py +193 -0
  73. prefect/blocks/core.py +3 -3
  74. prefect/blocks/notifications.py +8 -8
  75. prefect/client/base.py +1 -1
  76. prefect/client/cloud.py +1 -1
  77. prefect/client/orchestration.py +1 -1
  78. prefect/client/subscriptions.py +2 -6
  79. prefect/concurrency/services.py +1 -1
  80. prefect/context.py +3 -3
  81. prefect/deployments/deployments.py +3 -3
  82. prefect/engine.py +69 -9
  83. prefect/events/clients.py +1 -1
  84. prefect/filesystems.py +9 -9
  85. prefect/flow_runs.py +5 -1
  86. prefect/futures.py +1 -1
  87. prefect/infrastructure/container.py +3 -3
  88. prefect/infrastructure/kubernetes.py +4 -6
  89. prefect/infrastructure/process.py +3 -3
  90. prefect/input/run_input.py +1 -1
  91. prefect/logging/formatters.py +1 -1
  92. prefect/runner/server.py +3 -3
  93. prefect/settings.py +3 -4
  94. prefect/software/pip.py +1 -1
  95. prefect/task_engine.py +4 -0
  96. prefect/task_server.py +35 -17
  97. prefect/utilities/asyncutils.py +1 -1
  98. prefect/utilities/collections.py +1 -1
  99. {prefect_client-2.14.21.dist-info → prefect_client-2.15.0.dist-info}/METADATA +4 -2
  100. {prefect_client-2.14.21.dist-info → prefect_client-2.15.0.dist-info}/RECORD +103 -68
  101. {prefect_client-2.14.21.dist-info → prefect_client-2.15.0.dist-info}/LICENSE +0 -0
  102. {prefect_client-2.14.21.dist-info → prefect_client-2.15.0.dist-info}/WHEEL +0 -0
  103. {prefect_client-2.14.21.dist-info → prefect_client-2.15.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,805 @@
1
+ import contextlib
2
+ import inspect
3
+ import io
4
+ import json
5
+ import math
6
+ import queue
7
+ import typing
8
+ import warnings
9
+ from concurrent.futures import Future
10
+ from types import GeneratorType
11
+ from urllib.parse import unquote, urljoin
12
+
13
+ import anyio
14
+ import anyio.from_thread
15
+ from anyio.abc import ObjectReceiveStream, ObjectSendStream
16
+ from anyio.streams.stapled import StapledObjectStream
17
+ from prefect._vendor.starlette._utils import is_async_callable
18
+ from prefect._vendor.starlette.types import ASGIApp, Message, Receive, Scope, Send
19
+ from prefect._vendor.starlette.websockets import WebSocketDisconnect
20
+
21
+ try:
22
+ import httpx
23
+ except ModuleNotFoundError: # pragma: no cover
24
+ raise RuntimeError(
25
+ "The starlette.testclient module requires the httpx package to be installed.\n"
26
+ "You can install this with:\n"
27
+ " $ pip install httpx\n"
28
+ )
29
+ _PortalFactoryType = typing.Callable[
30
+ [], typing.ContextManager[anyio.abc.BlockingPortal]
31
+ ]
32
+
33
+ ASGIInstance = typing.Callable[[Receive, Send], typing.Awaitable[None]]
34
+ ASGI2App = typing.Callable[[Scope], ASGIInstance]
35
+ ASGI3App = typing.Callable[[Scope, Receive, Send], typing.Awaitable[None]]
36
+
37
+
38
+ _RequestData = typing.Mapping[str, typing.Union[str, typing.Iterable[str]]]
39
+
40
+
41
+ def _is_asgi3(app: typing.Union[ASGI2App, ASGI3App]) -> bool:
42
+ if inspect.isclass(app):
43
+ return hasattr(app, "__await__")
44
+ return is_async_callable(app)
45
+
46
+
47
+ class _WrapASGI2:
48
+ """
49
+ Provide an ASGI3 interface onto an ASGI2 app.
50
+ """
51
+
52
+ def __init__(self, app: ASGI2App) -> None:
53
+ self.app = app
54
+
55
+ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
56
+ instance = self.app(scope)
57
+ await instance(receive, send)
58
+
59
+
60
+ class _AsyncBackend(typing.TypedDict):
61
+ backend: str
62
+ backend_options: typing.Dict[str, typing.Any]
63
+
64
+
65
+ class _Upgrade(Exception):
66
+ def __init__(self, session: "WebSocketTestSession") -> None:
67
+ self.session = session
68
+
69
+
70
+ class WebSocketTestSession:
71
+ def __init__(
72
+ self,
73
+ app: ASGI3App,
74
+ scope: Scope,
75
+ portal_factory: _PortalFactoryType,
76
+ ) -> None:
77
+ self.app = app
78
+ self.scope = scope
79
+ self.accepted_subprotocol = None
80
+ self.portal_factory = portal_factory
81
+ self._receive_queue: "queue.Queue[Message]" = queue.Queue()
82
+ self._send_queue: "queue.Queue[Message | BaseException]" = queue.Queue()
83
+ self.extra_headers = None
84
+
85
+ def __enter__(self) -> "WebSocketTestSession":
86
+ self.exit_stack = contextlib.ExitStack()
87
+ self.portal = self.exit_stack.enter_context(self.portal_factory())
88
+
89
+ try:
90
+ _: "Future[None]" = self.portal.start_task_soon(self._run)
91
+ self.send({"type": "websocket.connect"})
92
+ message = self.receive()
93
+ self._raise_on_close(message)
94
+ except Exception:
95
+ self.exit_stack.close()
96
+ raise
97
+ self.accepted_subprotocol = message.get("subprotocol", None)
98
+ self.extra_headers = message.get("headers", None)
99
+ return self
100
+
101
+ def __exit__(self, *args: typing.Any) -> None:
102
+ try:
103
+ self.close(1000)
104
+ finally:
105
+ self.exit_stack.close()
106
+ while not self._send_queue.empty():
107
+ message = self._send_queue.get()
108
+ if isinstance(message, BaseException):
109
+ raise message
110
+
111
+ async def _run(self) -> None:
112
+ """
113
+ The sub-thread in which the websocket session runs.
114
+ """
115
+ scope = self.scope
116
+ receive = self._asgi_receive
117
+ send = self._asgi_send
118
+ try:
119
+ await self.app(scope, receive, send)
120
+ except BaseException as exc:
121
+ self._send_queue.put(exc)
122
+ raise
123
+
124
+ async def _asgi_receive(self) -> Message:
125
+ while self._receive_queue.empty():
126
+ await anyio.sleep(0)
127
+ return self._receive_queue.get()
128
+
129
+ async def _asgi_send(self, message: Message) -> None:
130
+ self._send_queue.put(message)
131
+
132
+ def _raise_on_close(self, message: Message) -> None:
133
+ if message["type"] == "websocket.close":
134
+ raise WebSocketDisconnect(
135
+ message.get("code", 1000), message.get("reason", "")
136
+ )
137
+
138
+ def send(self, message: Message) -> None:
139
+ self._receive_queue.put(message)
140
+
141
+ def send_text(self, data: str) -> None:
142
+ self.send({"type": "websocket.receive", "text": data})
143
+
144
+ def send_bytes(self, data: bytes) -> None:
145
+ self.send({"type": "websocket.receive", "bytes": data})
146
+
147
+ def send_json(self, data: typing.Any, mode: str = "text") -> None:
148
+ assert mode in ["text", "binary"]
149
+ text = json.dumps(data, separators=(",", ":"), ensure_ascii=False)
150
+ if mode == "text":
151
+ self.send({"type": "websocket.receive", "text": text})
152
+ else:
153
+ self.send({"type": "websocket.receive", "bytes": text.encode("utf-8")})
154
+
155
+ def close(self, code: int = 1000, reason: typing.Union[str, None] = None) -> None:
156
+ self.send({"type": "websocket.disconnect", "code": code, "reason": reason})
157
+
158
+ def receive(self) -> Message:
159
+ message = self._send_queue.get()
160
+ if isinstance(message, BaseException):
161
+ raise message
162
+ return message
163
+
164
+ def receive_text(self) -> str:
165
+ message = self.receive()
166
+ self._raise_on_close(message)
167
+ return typing.cast(str, message["text"])
168
+
169
+ def receive_bytes(self) -> bytes:
170
+ message = self.receive()
171
+ self._raise_on_close(message)
172
+ return typing.cast(bytes, message["bytes"])
173
+
174
+ def receive_json(self, mode: str = "text") -> typing.Any:
175
+ assert mode in ["text", "binary"]
176
+ message = self.receive()
177
+ self._raise_on_close(message)
178
+ if mode == "text":
179
+ text = message["text"]
180
+ else:
181
+ text = message["bytes"].decode("utf-8")
182
+ return json.loads(text)
183
+
184
+
185
+ class _TestClientTransport(httpx.BaseTransport):
186
+ def __init__(
187
+ self,
188
+ app: ASGI3App,
189
+ portal_factory: _PortalFactoryType,
190
+ raise_server_exceptions: bool = True,
191
+ root_path: str = "",
192
+ *,
193
+ app_state: typing.Dict[str, typing.Any],
194
+ ) -> None:
195
+ self.app = app
196
+ self.raise_server_exceptions = raise_server_exceptions
197
+ self.root_path = root_path
198
+ self.portal_factory = portal_factory
199
+ self.app_state = app_state
200
+
201
+ def handle_request(self, request: httpx.Request) -> httpx.Response:
202
+ scheme = request.url.scheme
203
+ netloc = request.url.netloc.decode(encoding="ascii")
204
+ path = request.url.path
205
+ raw_path = request.url.raw_path
206
+ query = request.url.query.decode(encoding="ascii")
207
+
208
+ default_port = {"http": 80, "ws": 80, "https": 443, "wss": 443}[scheme]
209
+
210
+ if ":" in netloc:
211
+ host, port_string = netloc.split(":", 1)
212
+ port = int(port_string)
213
+ else:
214
+ host = netloc
215
+ port = default_port
216
+
217
+ # Include the 'host' header.
218
+ if "host" in request.headers:
219
+ headers: typing.List[typing.Tuple[bytes, bytes]] = []
220
+ elif port == default_port: # pragma: no cover
221
+ headers = [(b"host", host.encode())]
222
+ else: # pragma: no cover
223
+ headers = [(b"host", (f"{host}:{port}").encode())]
224
+
225
+ # Include other request headers.
226
+ headers += [
227
+ (key.lower().encode(), value.encode())
228
+ for key, value in request.headers.multi_items()
229
+ ]
230
+
231
+ scope: typing.Dict[str, typing.Any]
232
+
233
+ if scheme in {"ws", "wss"}:
234
+ subprotocol = request.headers.get("sec-websocket-protocol", None)
235
+ if subprotocol is None:
236
+ subprotocols: typing.Sequence[str] = []
237
+ else:
238
+ subprotocols = [value.strip() for value in subprotocol.split(",")]
239
+ scope = {
240
+ "type": "websocket",
241
+ "path": unquote(path),
242
+ "raw_path": raw_path,
243
+ "root_path": self.root_path,
244
+ "scheme": scheme,
245
+ "query_string": query.encode(),
246
+ "headers": headers,
247
+ "client": ["testclient", 50000],
248
+ "server": [host, port],
249
+ "subprotocols": subprotocols,
250
+ "state": self.app_state.copy(),
251
+ }
252
+ session = WebSocketTestSession(self.app, scope, self.portal_factory)
253
+ raise _Upgrade(session)
254
+
255
+ scope = {
256
+ "type": "http",
257
+ "http_version": "1.1",
258
+ "method": request.method,
259
+ "path": unquote(path),
260
+ "raw_path": raw_path,
261
+ "root_path": self.root_path,
262
+ "scheme": scheme,
263
+ "query_string": query.encode(),
264
+ "headers": headers,
265
+ "client": ["testclient", 50000],
266
+ "server": [host, port],
267
+ "extensions": {"http.response.debug": {}},
268
+ "state": self.app_state.copy(),
269
+ }
270
+
271
+ request_complete = False
272
+ response_started = False
273
+ response_complete: anyio.Event
274
+ raw_kwargs: typing.Dict[str, typing.Any] = {"stream": io.BytesIO()}
275
+ template = None
276
+ context = None
277
+
278
+ async def receive() -> Message:
279
+ nonlocal request_complete
280
+
281
+ if request_complete:
282
+ if not response_complete.is_set():
283
+ await response_complete.wait()
284
+ return {"type": "http.disconnect"}
285
+
286
+ body = request.read()
287
+ if isinstance(body, str):
288
+ body_bytes: bytes = body.encode("utf-8") # pragma: no cover
289
+ elif body is None:
290
+ body_bytes = b"" # pragma: no cover
291
+ elif isinstance(body, GeneratorType):
292
+ try: # pragma: no cover
293
+ chunk = body.send(None)
294
+ if isinstance(chunk, str):
295
+ chunk = chunk.encode("utf-8")
296
+ return {"type": "http.request", "body": chunk, "more_body": True}
297
+ except StopIteration: # pragma: no cover
298
+ request_complete = True
299
+ return {"type": "http.request", "body": b""}
300
+ else:
301
+ body_bytes = body
302
+
303
+ request_complete = True
304
+ return {"type": "http.request", "body": body_bytes}
305
+
306
+ async def send(message: Message) -> None:
307
+ nonlocal raw_kwargs, response_started, template, context
308
+
309
+ if message["type"] == "http.response.start":
310
+ assert (
311
+ not response_started
312
+ ), 'Received multiple "http.response.start" messages.'
313
+ raw_kwargs["status_code"] = message["status"]
314
+ raw_kwargs["headers"] = [
315
+ (key.decode(), value.decode())
316
+ for key, value in message.get("headers", [])
317
+ ]
318
+ response_started = True
319
+ elif message["type"] == "http.response.body":
320
+ assert (
321
+ response_started
322
+ ), 'Received "http.response.body" without "http.response.start".'
323
+ assert (
324
+ not response_complete.is_set()
325
+ ), 'Received "http.response.body" after response completed.'
326
+ body = message.get("body", b"")
327
+ more_body = message.get("more_body", False)
328
+ if request.method != "HEAD":
329
+ raw_kwargs["stream"].write(body)
330
+ if not more_body:
331
+ raw_kwargs["stream"].seek(0)
332
+ response_complete.set()
333
+ elif message["type"] == "http.response.debug":
334
+ template = message["info"]["template"]
335
+ context = message["info"]["context"]
336
+
337
+ try:
338
+ with self.portal_factory() as portal:
339
+ response_complete = portal.call(anyio.Event)
340
+ portal.call(self.app, scope, receive, send)
341
+ except BaseException as exc:
342
+ if self.raise_server_exceptions:
343
+ raise exc
344
+
345
+ if self.raise_server_exceptions:
346
+ assert response_started, "TestClient did not receive any response."
347
+ elif not response_started:
348
+ raw_kwargs = {
349
+ "status_code": 500,
350
+ "headers": [],
351
+ "stream": io.BytesIO(),
352
+ }
353
+
354
+ raw_kwargs["stream"] = httpx.ByteStream(raw_kwargs["stream"].read())
355
+
356
+ response = httpx.Response(**raw_kwargs, request=request)
357
+ if template is not None:
358
+ response.template = template # type: ignore[attr-defined]
359
+ response.context = context # type: ignore[attr-defined]
360
+ return response
361
+
362
+
363
+ class TestClient(httpx.Client):
364
+ __test__ = False
365
+ task: "Future[None]"
366
+ portal: typing.Optional[anyio.abc.BlockingPortal] = None
367
+
368
+ def __init__(
369
+ self,
370
+ app: ASGIApp,
371
+ base_url: str = "http://testserver",
372
+ raise_server_exceptions: bool = True,
373
+ root_path: str = "",
374
+ backend: str = "asyncio",
375
+ backend_options: typing.Optional[typing.Dict[str, typing.Any]] = None,
376
+ cookies: httpx._types.CookieTypes = None,
377
+ headers: typing.Dict[str, str] = None,
378
+ follow_redirects: bool = True,
379
+ ) -> None:
380
+ self.async_backend = _AsyncBackend(
381
+ backend=backend, backend_options=backend_options or {}
382
+ )
383
+ if _is_asgi3(app):
384
+ app = typing.cast(ASGI3App, app)
385
+ asgi_app = app
386
+ else:
387
+ app = typing.cast(ASGI2App, app) # type: ignore[assignment]
388
+ asgi_app = _WrapASGI2(app) # type: ignore[arg-type]
389
+ self.app = asgi_app
390
+ self.app_state: typing.Dict[str, typing.Any] = {}
391
+ transport = _TestClientTransport(
392
+ self.app,
393
+ portal_factory=self._portal_factory,
394
+ raise_server_exceptions=raise_server_exceptions,
395
+ root_path=root_path,
396
+ app_state=self.app_state,
397
+ )
398
+ if headers is None:
399
+ headers = {}
400
+ headers.setdefault("user-agent", "testclient")
401
+ super().__init__(
402
+ app=self.app,
403
+ base_url=base_url,
404
+ headers=headers,
405
+ transport=transport,
406
+ follow_redirects=follow_redirects,
407
+ cookies=cookies,
408
+ )
409
+
410
+ @contextlib.contextmanager
411
+ def _portal_factory(self) -> typing.Generator[anyio.abc.BlockingPortal, None, None]:
412
+ if self.portal is not None:
413
+ yield self.portal
414
+ else:
415
+ with anyio.from_thread.start_blocking_portal(
416
+ **self.async_backend
417
+ ) as portal:
418
+ yield portal
419
+
420
+ def _choose_redirect_arg(
421
+ self,
422
+ follow_redirects: typing.Optional[bool],
423
+ allow_redirects: typing.Optional[bool],
424
+ ) -> typing.Union[bool, httpx._client.UseClientDefault]:
425
+ redirect: typing.Union[
426
+ bool, httpx._client.UseClientDefault
427
+ ] = httpx._client.USE_CLIENT_DEFAULT
428
+ if allow_redirects is not None:
429
+ message = (
430
+ "The `allow_redirects` argument is deprecated. "
431
+ "Use `follow_redirects` instead."
432
+ )
433
+ warnings.warn(message, DeprecationWarning)
434
+ redirect = allow_redirects
435
+ if follow_redirects is not None:
436
+ redirect = follow_redirects
437
+ elif allow_redirects is not None and follow_redirects is not None:
438
+ raise RuntimeError( # pragma: no cover
439
+ "Cannot use both `allow_redirects` and `follow_redirects`."
440
+ )
441
+ return redirect
442
+
443
+ def request( # type: ignore[override]
444
+ self,
445
+ method: str,
446
+ url: httpx._types.URLTypes,
447
+ *,
448
+ content: typing.Optional[httpx._types.RequestContent] = None,
449
+ data: typing.Optional[_RequestData] = None,
450
+ files: typing.Optional[httpx._types.RequestFiles] = None,
451
+ json: typing.Any = None,
452
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
453
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
454
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
455
+ auth: typing.Union[
456
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
457
+ ] = httpx._client.USE_CLIENT_DEFAULT,
458
+ follow_redirects: typing.Optional[bool] = None,
459
+ allow_redirects: typing.Optional[bool] = None,
460
+ timeout: typing.Union[
461
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
462
+ ] = httpx._client.USE_CLIENT_DEFAULT,
463
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
464
+ ) -> httpx.Response:
465
+ url = self.base_url.join(url)
466
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
467
+ return super().request(
468
+ method,
469
+ url,
470
+ content=content,
471
+ data=data,
472
+ files=files,
473
+ json=json,
474
+ params=params,
475
+ headers=headers,
476
+ cookies=cookies,
477
+ auth=auth,
478
+ follow_redirects=redirect,
479
+ timeout=timeout,
480
+ extensions=extensions,
481
+ )
482
+
483
+ def get( # type: ignore[override]
484
+ self,
485
+ url: httpx._types.URLTypes,
486
+ *,
487
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
488
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
489
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
490
+ auth: typing.Union[
491
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
492
+ ] = httpx._client.USE_CLIENT_DEFAULT,
493
+ follow_redirects: typing.Optional[bool] = None,
494
+ allow_redirects: typing.Optional[bool] = None,
495
+ timeout: typing.Union[
496
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
497
+ ] = httpx._client.USE_CLIENT_DEFAULT,
498
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
499
+ ) -> httpx.Response:
500
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
501
+ return super().get(
502
+ url,
503
+ params=params,
504
+ headers=headers,
505
+ cookies=cookies,
506
+ auth=auth,
507
+ follow_redirects=redirect,
508
+ timeout=timeout,
509
+ extensions=extensions,
510
+ )
511
+
512
+ def options( # type: ignore[override]
513
+ self,
514
+ url: httpx._types.URLTypes,
515
+ *,
516
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
517
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
518
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
519
+ auth: typing.Union[
520
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
521
+ ] = httpx._client.USE_CLIENT_DEFAULT,
522
+ follow_redirects: typing.Optional[bool] = None,
523
+ allow_redirects: typing.Optional[bool] = None,
524
+ timeout: typing.Union[
525
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
526
+ ] = httpx._client.USE_CLIENT_DEFAULT,
527
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
528
+ ) -> httpx.Response:
529
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
530
+ return super().options(
531
+ url,
532
+ params=params,
533
+ headers=headers,
534
+ cookies=cookies,
535
+ auth=auth,
536
+ follow_redirects=redirect,
537
+ timeout=timeout,
538
+ extensions=extensions,
539
+ )
540
+
541
+ def head( # type: ignore[override]
542
+ self,
543
+ url: httpx._types.URLTypes,
544
+ *,
545
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
546
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
547
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
548
+ auth: typing.Union[
549
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
550
+ ] = httpx._client.USE_CLIENT_DEFAULT,
551
+ follow_redirects: typing.Optional[bool] = None,
552
+ allow_redirects: typing.Optional[bool] = None,
553
+ timeout: typing.Union[
554
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
555
+ ] = httpx._client.USE_CLIENT_DEFAULT,
556
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
557
+ ) -> httpx.Response:
558
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
559
+ return super().head(
560
+ url,
561
+ params=params,
562
+ headers=headers,
563
+ cookies=cookies,
564
+ auth=auth,
565
+ follow_redirects=redirect,
566
+ timeout=timeout,
567
+ extensions=extensions,
568
+ )
569
+
570
+ def post( # type: ignore[override]
571
+ self,
572
+ url: httpx._types.URLTypes,
573
+ *,
574
+ content: typing.Optional[httpx._types.RequestContent] = None,
575
+ data: typing.Optional[_RequestData] = None,
576
+ files: typing.Optional[httpx._types.RequestFiles] = None,
577
+ json: typing.Any = None,
578
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
579
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
580
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
581
+ auth: typing.Union[
582
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
583
+ ] = httpx._client.USE_CLIENT_DEFAULT,
584
+ follow_redirects: typing.Optional[bool] = None,
585
+ allow_redirects: typing.Optional[bool] = None,
586
+ timeout: typing.Union[
587
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
588
+ ] = httpx._client.USE_CLIENT_DEFAULT,
589
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
590
+ ) -> httpx.Response:
591
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
592
+ return super().post(
593
+ url,
594
+ content=content,
595
+ data=data,
596
+ files=files,
597
+ json=json,
598
+ params=params,
599
+ headers=headers,
600
+ cookies=cookies,
601
+ auth=auth,
602
+ follow_redirects=redirect,
603
+ timeout=timeout,
604
+ extensions=extensions,
605
+ )
606
+
607
+ def put( # type: ignore[override]
608
+ self,
609
+ url: httpx._types.URLTypes,
610
+ *,
611
+ content: typing.Optional[httpx._types.RequestContent] = None,
612
+ data: typing.Optional[_RequestData] = None,
613
+ files: typing.Optional[httpx._types.RequestFiles] = None,
614
+ json: typing.Any = None,
615
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
616
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
617
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
618
+ auth: typing.Union[
619
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
620
+ ] = httpx._client.USE_CLIENT_DEFAULT,
621
+ follow_redirects: typing.Optional[bool] = None,
622
+ allow_redirects: typing.Optional[bool] = None,
623
+ timeout: typing.Union[
624
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
625
+ ] = httpx._client.USE_CLIENT_DEFAULT,
626
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
627
+ ) -> httpx.Response:
628
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
629
+ return super().put(
630
+ url,
631
+ content=content,
632
+ data=data,
633
+ files=files,
634
+ json=json,
635
+ params=params,
636
+ headers=headers,
637
+ cookies=cookies,
638
+ auth=auth,
639
+ follow_redirects=redirect,
640
+ timeout=timeout,
641
+ extensions=extensions,
642
+ )
643
+
644
+ def patch( # type: ignore[override]
645
+ self,
646
+ url: httpx._types.URLTypes,
647
+ *,
648
+ content: typing.Optional[httpx._types.RequestContent] = None,
649
+ data: typing.Optional[_RequestData] = None,
650
+ files: typing.Optional[httpx._types.RequestFiles] = None,
651
+ json: typing.Any = None,
652
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
653
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
654
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
655
+ auth: typing.Union[
656
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
657
+ ] = httpx._client.USE_CLIENT_DEFAULT,
658
+ follow_redirects: typing.Optional[bool] = None,
659
+ allow_redirects: typing.Optional[bool] = None,
660
+ timeout: typing.Union[
661
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
662
+ ] = httpx._client.USE_CLIENT_DEFAULT,
663
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
664
+ ) -> httpx.Response:
665
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
666
+ return super().patch(
667
+ url,
668
+ content=content,
669
+ data=data,
670
+ files=files,
671
+ json=json,
672
+ params=params,
673
+ headers=headers,
674
+ cookies=cookies,
675
+ auth=auth,
676
+ follow_redirects=redirect,
677
+ timeout=timeout,
678
+ extensions=extensions,
679
+ )
680
+
681
+ def delete( # type: ignore[override]
682
+ self,
683
+ url: httpx._types.URLTypes,
684
+ *,
685
+ params: typing.Optional[httpx._types.QueryParamTypes] = None,
686
+ headers: typing.Optional[httpx._types.HeaderTypes] = None,
687
+ cookies: typing.Optional[httpx._types.CookieTypes] = None,
688
+ auth: typing.Union[
689
+ httpx._types.AuthTypes, httpx._client.UseClientDefault
690
+ ] = httpx._client.USE_CLIENT_DEFAULT,
691
+ follow_redirects: typing.Optional[bool] = None,
692
+ allow_redirects: typing.Optional[bool] = None,
693
+ timeout: typing.Union[
694
+ httpx._types.TimeoutTypes, httpx._client.UseClientDefault
695
+ ] = httpx._client.USE_CLIENT_DEFAULT,
696
+ extensions: typing.Optional[typing.Dict[str, typing.Any]] = None,
697
+ ) -> httpx.Response:
698
+ redirect = self._choose_redirect_arg(follow_redirects, allow_redirects)
699
+ return super().delete(
700
+ url,
701
+ params=params,
702
+ headers=headers,
703
+ cookies=cookies,
704
+ auth=auth,
705
+ follow_redirects=redirect,
706
+ timeout=timeout,
707
+ extensions=extensions,
708
+ )
709
+
710
+ def websocket_connect(
711
+ self, url: str, subprotocols: typing.Sequence[str] = None, **kwargs: typing.Any
712
+ ) -> "WebSocketTestSession":
713
+ url = urljoin("ws://testserver", url)
714
+ headers = kwargs.get("headers", {})
715
+ headers.setdefault("connection", "upgrade")
716
+ headers.setdefault("sec-websocket-key", "testserver==")
717
+ headers.setdefault("sec-websocket-version", "13")
718
+ if subprotocols is not None:
719
+ headers.setdefault("sec-websocket-protocol", ", ".join(subprotocols))
720
+ kwargs["headers"] = headers
721
+ try:
722
+ super().request("GET", url, **kwargs)
723
+ except _Upgrade as exc:
724
+ session = exc.session
725
+ else:
726
+ raise RuntimeError("Expected WebSocket upgrade") # pragma: no cover
727
+
728
+ return session
729
+
730
+ def __enter__(self) -> "TestClient":
731
+ with contextlib.ExitStack() as stack:
732
+ self.portal = portal = stack.enter_context(
733
+ anyio.from_thread.start_blocking_portal(**self.async_backend)
734
+ )
735
+
736
+ @stack.callback
737
+ def reset_portal() -> None:
738
+ self.portal = None
739
+
740
+ send1: ObjectSendStream[
741
+ typing.Optional[typing.MutableMapping[str, typing.Any]]
742
+ ]
743
+ receive1: ObjectReceiveStream[
744
+ typing.Optional[typing.MutableMapping[str, typing.Any]]
745
+ ]
746
+ send2: ObjectSendStream[typing.MutableMapping[str, typing.Any]]
747
+ receive2: ObjectReceiveStream[typing.MutableMapping[str, typing.Any]]
748
+ send1, receive1 = anyio.create_memory_object_stream(math.inf)
749
+ send2, receive2 = anyio.create_memory_object_stream(math.inf)
750
+ self.stream_send = StapledObjectStream(send1, receive1)
751
+ self.stream_receive = StapledObjectStream(send2, receive2)
752
+ self.task = portal.start_task_soon(self.lifespan)
753
+ portal.call(self.wait_startup)
754
+
755
+ @stack.callback
756
+ def wait_shutdown() -> None:
757
+ portal.call(self.wait_shutdown)
758
+
759
+ self.exit_stack = stack.pop_all()
760
+
761
+ return self
762
+
763
+ def __exit__(self, *args: typing.Any) -> None:
764
+ self.exit_stack.close()
765
+
766
+ async def lifespan(self) -> None:
767
+ scope = {"type": "lifespan", "state": self.app_state}
768
+ try:
769
+ await self.app(scope, self.stream_receive.receive, self.stream_send.send)
770
+ finally:
771
+ await self.stream_send.send(None)
772
+
773
+ async def wait_startup(self) -> None:
774
+ await self.stream_receive.send({"type": "lifespan.startup"})
775
+
776
+ async def receive() -> typing.Any:
777
+ message = await self.stream_send.receive()
778
+ if message is None:
779
+ self.task.result()
780
+ return message
781
+
782
+ message = await receive()
783
+ assert message["type"] in (
784
+ "lifespan.startup.complete",
785
+ "lifespan.startup.failed",
786
+ )
787
+ if message["type"] == "lifespan.startup.failed":
788
+ await receive()
789
+
790
+ async def wait_shutdown(self) -> None:
791
+ async def receive() -> typing.Any:
792
+ message = await self.stream_send.receive()
793
+ if message is None:
794
+ self.task.result()
795
+ return message
796
+
797
+ async with self.stream_send:
798
+ await self.stream_receive.send({"type": "lifespan.shutdown"})
799
+ message = await receive()
800
+ assert message["type"] in (
801
+ "lifespan.shutdown.complete",
802
+ "lifespan.shutdown.failed",
803
+ )
804
+ if message["type"] == "lifespan.shutdown.failed":
805
+ await receive()