wehub-resource-sync
92 lines
2.5 KiB
Python
92 lines
2.5 KiB
Python
import asyncio
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from ray import serve
|
|
from ray._common.test_utils import SignalActor
|
|
from ray.serve._private.constants import (
|
|
SERVE_HTTP_REQUEST_DISCONNECT_DISABLED_HEADER,
|
|
SERVE_HTTP_REQUEST_TIMEOUT_S_HEADER,
|
|
)
|
|
from ray.serve._private.test_utils import get_application_url
|
|
from ray.serve.config import HTTPOptions
|
|
from ray.serve.tests.conftest import * # noqa
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("set_http_options_timeout", [False, True])
|
|
@pytest.mark.parametrize("timeout_header", [None, "abc", "0", "1"])
|
|
@pytest.mark.parametrize("disconnect_header", [None, "?0", "?1"])
|
|
async def test_http_request_timeout_disconnect_headers(
|
|
ray_instance,
|
|
ray_shutdown,
|
|
set_http_options_timeout,
|
|
disconnect_header,
|
|
timeout_header,
|
|
):
|
|
"""Test cases of request timeout and disconnect header."""
|
|
http_request_timeout_s = None
|
|
if set_http_options_timeout:
|
|
http_request_timeout_s = 1
|
|
|
|
serve.start(
|
|
http_options=HTTPOptions(
|
|
host="0.0.0.0",
|
|
request_timeout_s=http_request_timeout_s,
|
|
),
|
|
)
|
|
|
|
signal = SignalActor.remote()
|
|
|
|
@serve.deployment
|
|
class Model:
|
|
async def __call__(self):
|
|
await signal.wait.remote()
|
|
|
|
serve.run(Model.bind())
|
|
|
|
http_url = get_application_url("HTTP")
|
|
|
|
headers = {}
|
|
if timeout_header is not None:
|
|
headers[
|
|
SERVE_HTTP_REQUEST_TIMEOUT_S_HEADER.encode("utf-8")
|
|
] = timeout_header.encode("utf-8")
|
|
if disconnect_header is not None:
|
|
headers[
|
|
SERVE_HTTP_REQUEST_DISCONNECT_DISABLED_HEADER.encode("utf-8")
|
|
] = disconnect_header.encode("utf-8")
|
|
|
|
expect_timeout = (
|
|
timeout_header == "1"
|
|
or set_http_options_timeout is True
|
|
and timeout_header != "0"
|
|
)
|
|
expect_disconnect = not expect_timeout and disconnect_header != "?1"
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
response_task = asyncio.create_task(
|
|
client.get(http_url, headers=headers, timeout=2)
|
|
)
|
|
await asyncio.sleep(0.01)
|
|
|
|
if not expect_timeout and not expect_disconnect:
|
|
signal.send.remote()
|
|
|
|
if expect_disconnect:
|
|
with pytest.raises(httpx.ReadTimeout):
|
|
await response_task
|
|
else:
|
|
response = await response_task
|
|
if expect_timeout:
|
|
assert response.status_code == 408
|
|
else:
|
|
assert response.status_code == 200
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|