chore: import upstream snapshot with attribution

python/ray/serve/tests/test_per_request_headers.py

@@ -0,0 +1,91 @@
+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__]))