Files
ray-project--ray/python/ray/serve/tests/unit/test_proxy_router.py
T
2026-07-13 13:17:40 +08:00

228 lines
7.6 KiB
Python

import pytest
from ray.serve._private.common import DeploymentID, EndpointInfo
from ray.serve._private.proxy_router import (
NO_REPLICAS_MESSAGE,
NO_ROUTES_MESSAGE,
ProxyRouter,
)
from ray.serve._private.test_utils import MockDeploymentHandle
@pytest.fixture
def mock_router():
def mock_get_handle(endpoint, info):
return MockDeploymentHandle(endpoint.name, endpoint.app_name)
yield ProxyRouter(mock_get_handle)
def test_no_match(mock_router: ProxyRouter):
router = mock_router
router.update_routes(
{
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
route="/hello"
),
DeploymentID(name="endpoint2", app_name="app2"): EndpointInfo(
route="/hello2"
),
},
)
assert router.match_route("/nonexistent") is None
assert router.get_handle_for_endpoint("/nonexistent") is None
def test_default_route(mock_router: ProxyRouter):
router = mock_router
router.update_routes(
{
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
route="/endpoint"
),
DeploymentID(name="endpoint2", app_name="app2"): EndpointInfo(
route="/endpoint2"
),
},
)
# Route based matching
assert router.match_route("/nonexistent") is None
route, handle, app_is_cross_language = router.match_route("/endpoint")
assert route == "/endpoint"
assert handle == ("endpoint", "default")
assert not app_is_cross_language
# Endpoint based matching
assert router.get_handle_for_endpoint("/nonexistent") is None
route, handle, app_is_cross_language = router.get_handle_for_endpoint("default")
assert route == "/endpoint"
assert handle == ("endpoint", "default")
assert not app_is_cross_language
def test_trailing_slash(mock_router: ProxyRouter):
router = mock_router
router.update_routes(
{
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
route="/test"
)
},
)
route, handle, _ = router.match_route("/test/")
assert route == "/test" and handle == ("endpoint", "default")
router.update_routes(
{
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
route="/test/"
)
},
)
assert router.match_route("/test") is None
def test_prefix_match(mock_router):
router = mock_router
router.update_routes(
{
DeploymentID(name="endpoint1", app_name="default"): EndpointInfo(
route="/test/test2"
),
DeploymentID(name="endpoint2", app_name="default"): EndpointInfo(
route="/test"
),
DeploymentID(name="endpoint3", app_name="default"): EndpointInfo(route="/"),
},
)
route, handle, _ = router.match_route("/test/test2/subpath")
assert route == "/test/test2" and handle == ("endpoint1", "default")
route, handle, _ = router.match_route("/test/test2/")
assert route == "/test/test2" and handle == ("endpoint1", "default")
route, handle, _ = router.match_route("/test/test2")
assert route == "/test/test2" and handle == ("endpoint1", "default")
route, handle, _ = router.match_route("/test/subpath")
assert route == "/test" and handle == ("endpoint2", "default")
route, handle, _ = router.match_route("/test/")
assert route == "/test" and handle == ("endpoint2", "default")
route, handle, _ = router.match_route("/test")
assert route == "/test" and handle == ("endpoint2", "default")
route, handle, _ = router.match_route("/test2")
assert route == "/" and handle == ("endpoint3", "default")
route, handle, _ = router.match_route("/")
assert route == "/" and handle == ("endpoint3", "default")
def test_update_routes(mock_router):
router = mock_router
router.update_routes(
{DeploymentID("endpoint", "app1"): EndpointInfo(route="/endpoint")},
)
route, handle, app_is_cross_language = router.match_route("/endpoint")
assert route == "/endpoint"
assert handle == ("endpoint", "app1")
assert not app_is_cross_language
route, handle, app_is_cross_language = router.get_handle_for_endpoint("app1")
assert route == "/endpoint"
assert handle == ("endpoint", "app1")
assert not app_is_cross_language
router.update_routes(
{
DeploymentID(name="endpoint2", app_name="app2"): EndpointInfo(
route="/endpoint2",
app_is_cross_language=True,
),
DeploymentID(name="endpoint3", app_name="app3"): EndpointInfo(
route="/endpoint3",
app_is_cross_language=True,
),
},
)
assert router.match_route("/endpoint") is None
assert router.match_route("app1") is None
route, handle, app_is_cross_language = router.match_route("/endpoint2")
assert route == "/endpoint2"
assert handle == ("endpoint2", "app2")
assert app_is_cross_language
route, handle, app_is_cross_language = router.get_handle_for_endpoint("app2")
assert route == "/endpoint2"
assert handle == ("endpoint2", "app2")
assert app_is_cross_language
class TestReadyForTraffic:
@pytest.mark.parametrize("is_head", [False, True])
def test_route_table_not_populated(self, mock_router, is_head: bool):
"""Proxy router should NOT be ready for traffic if:
- it has not received route table from controller
"""
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=is_head)
assert not ready_for_traffic
assert msg == NO_ROUTES_MESSAGE
def test_head_route_table_populated_no_replicas(self, mock_router):
"""Proxy router should be ready for traffic if:
- it has received route table from controller
- it hasn't received any replicas yet
- it lives on head node
"""
d_id = DeploymentID(name="A", app_name="B")
mock_router.update_routes({d_id: EndpointInfo(route="/")})
mock_router.handles[d_id].set_running_replicas_populated(False)
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=True)
assert ready_for_traffic
assert not msg
def test_worker_route_table_populated_no_replicas(self, mock_router):
"""Proxy router should NOT be ready for traffic if:
- it has received route table from controller
- it hasn't received any replicas yet
- it lives on a worker node
"""
d_id = DeploymentID(name="A", app_name="B")
mock_router.update_routes({d_id: EndpointInfo(route="/")})
mock_router.handles[d_id].set_running_replicas_populated(False)
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=False)
assert not ready_for_traffic
assert msg == NO_REPLICAS_MESSAGE
@pytest.mark.parametrize("is_head", [False, True])
def test_route_table_populated_with_replicas(self, mock_router, is_head: bool):
"""Proxy router should be ready for traffic if:
- it has received route table from controller
- it has received replicas from controller
"""
d_id = DeploymentID(name="A", app_name="B")
mock_router.update_routes({d_id: EndpointInfo(route="/")})
mock_router.handles[d_id].set_running_replicas_populated(True)
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=is_head)
assert ready_for_traffic
assert not msg
if __name__ == "__main__":
import sys
sys.exit(pytest.main(["-v", "-s", __file__]))