wehub-resource-sync
225 lines
7.5 KiB
Python
225 lines
7.5 KiB
Python
import pickle
|
|
from typing import Callable
|
|
from unittest.mock import Mock
|
|
|
|
import grpc
|
|
import pytest
|
|
from google.protobuf.any_pb2 import Any as AnyProto
|
|
|
|
from ray import cloudpickle
|
|
from ray.serve._private.default_impl import add_grpc_address
|
|
from ray.serve._private.grpc_util import (
|
|
GRPC_MAX_STATUS_DETAILS_LENGTH,
|
|
_truncate_message,
|
|
get_grpc_response_status,
|
|
gRPCGenericServer,
|
|
)
|
|
from ray.serve._private.proxy_request_response import gRPCStreamingType
|
|
from ray.serve._private.test_utils import FakeGrpcContext
|
|
from ray.serve.exceptions import BackPressureError, gRPCStatusError
|
|
from ray.serve.grpc_util import RayServegRPCContext
|
|
|
|
|
|
class FakeGrpcServer:
|
|
def __init__(self):
|
|
self.address = None
|
|
|
|
def add_insecure_port(self, address):
|
|
self.address = address
|
|
|
|
|
|
def fake_service_handler_factory(
|
|
service_method: str, streaming_type: gRPCStreamingType
|
|
) -> Callable:
|
|
def foo() -> bytes:
|
|
return f"{streaming_type.value} call from {service_method}".encode()
|
|
|
|
return foo
|
|
|
|
|
|
def test_grpc_server():
|
|
"""Test `gRPCGenericServer` did the correct overrides.
|
|
|
|
When a add_servicer_to_server function is called on an instance of `gRPCGenericServer`,
|
|
it correctly overrides `response_serializer` to None, and `unary_unary`,
|
|
`unary_stream`, `stream_unary`, and `stream_stream` to be generated from the
|
|
factory function.
|
|
"""
|
|
service_name = "ray.serve.ServeAPIService"
|
|
method_name = "ServeRoutes"
|
|
|
|
def add_test_servicer_to_server(servicer, server):
|
|
rpc_method_handlers = {
|
|
method_name: grpc.unary_unary_rpc_method_handler(
|
|
servicer.ServeRoutes,
|
|
request_deserializer=AnyProto.FromString,
|
|
response_serializer=AnyProto.SerializeToString,
|
|
),
|
|
}
|
|
generic_handler = grpc.method_handlers_generic_handler(
|
|
service_name, rpc_method_handlers
|
|
)
|
|
server.add_generic_rpc_handlers((generic_handler,))
|
|
|
|
grpc_server = gRPCGenericServer(fake_service_handler_factory)
|
|
dummy_servicer = Mock()
|
|
|
|
# Ensure `generic_rpc_handlers` is not populated before calling
|
|
# the add_servicer_to_server function.
|
|
assert grpc_server.generic_rpc_handlers == []
|
|
|
|
add_test_servicer_to_server(dummy_servicer, grpc_server)
|
|
|
|
# `generic_rpc_handlers` should be populated after add_servicer_to_server is called.
|
|
assert len(grpc_server.generic_rpc_handlers) == 1
|
|
|
|
# The populated rpc handler should have the correct service name.
|
|
rpc_handler = grpc_server.generic_rpc_handlers[0][0]
|
|
assert rpc_handler.service_name() == service_name
|
|
|
|
# The populated method handlers should have the correct response_serializer,
|
|
# unary_unary, unary_stream, stream_unary, and stream_stream.
|
|
service_method = f"/{service_name}/{method_name}"
|
|
method_handlers = rpc_handler._method_handlers.get(service_method)
|
|
assert method_handlers.response_serializer is None
|
|
assert (
|
|
method_handlers.unary_unary()
|
|
== f"unary_unary call from {service_method}".encode()
|
|
)
|
|
assert (
|
|
method_handlers.unary_stream()
|
|
== f"unary_stream call from {service_method}".encode()
|
|
)
|
|
assert (
|
|
method_handlers.stream_unary()
|
|
== f"stream_unary call from {service_method}".encode()
|
|
)
|
|
assert (
|
|
method_handlers.stream_stream()
|
|
== f"stream_stream call from {service_method}".encode()
|
|
)
|
|
|
|
|
|
def test_ray_serve_grpc_context_serializable():
|
|
"""RayServegRPCContext should be serializable."""
|
|
context = RayServegRPCContext(FakeGrpcContext())
|
|
pickled_context = pickle.dumps(context)
|
|
deserialized_context = pickle.loads(pickled_context)
|
|
assert deserialized_context.__dict__ == context.__dict__
|
|
|
|
cloudpickled_context = cloudpickle.dumps(context)
|
|
deserialized_context = pickle.loads(cloudpickled_context)
|
|
assert deserialized_context.__dict__ == context.__dict__
|
|
|
|
|
|
def test_add_grpc_address():
|
|
"""Test `add_grpc_address` adds the address to the gRPC server."""
|
|
fake_grpc_server = FakeGrpcServer()
|
|
grpc_address = "fake_address:50051"
|
|
assert fake_grpc_server.address is None
|
|
add_grpc_address(fake_grpc_server, grpc_address)
|
|
assert fake_grpc_server.address == grpc_address
|
|
|
|
|
|
def test_get_grpc_response_status_backpressure_error():
|
|
"""Test that BackPressureError returns RESOURCE_EXHAUSTED status."""
|
|
backpressure_error = BackPressureError(
|
|
num_queued_requests=10, max_queued_requests=5
|
|
)
|
|
|
|
status = get_grpc_response_status(
|
|
exc=backpressure_error, request_timeout_s=30.0, request_id="test_request_123"
|
|
)
|
|
|
|
assert status.code == grpc.StatusCode.RESOURCE_EXHAUSTED
|
|
assert status.is_error is True
|
|
assert status.message == backpressure_error.message
|
|
|
|
|
|
def test_get_grpc_response_status_grpc_status_error():
|
|
"""Test that gRPCStatusError preserves user-set status code."""
|
|
original_error = RuntimeError("test error")
|
|
user_status_code = grpc.StatusCode.INVALID_ARGUMENT
|
|
user_details = "Invalid argument provided"
|
|
|
|
grpc_status_error = gRPCStatusError(
|
|
original_exception=original_error,
|
|
code=user_status_code,
|
|
details=user_details,
|
|
)
|
|
|
|
status = get_grpc_response_status(
|
|
exc=grpc_status_error, request_timeout_s=30.0, request_id="test_request_123"
|
|
)
|
|
|
|
assert status.code == user_status_code
|
|
assert status.is_error is True
|
|
assert status.message == user_details
|
|
|
|
|
|
def test_get_grpc_response_status_grpc_status_error_no_details():
|
|
"""Test that gRPCStatusError without details uses original exception message."""
|
|
original_error = RuntimeError("original error message")
|
|
user_status_code = grpc.StatusCode.RESOURCE_EXHAUSTED
|
|
|
|
grpc_status_error = gRPCStatusError(
|
|
original_exception=original_error,
|
|
code=user_status_code,
|
|
details=None,
|
|
)
|
|
|
|
status = get_grpc_response_status(
|
|
exc=grpc_status_error, request_timeout_s=30.0, request_id="test_request_123"
|
|
)
|
|
|
|
assert status.code == user_status_code
|
|
assert status.is_error is True
|
|
assert "original error message" in status.message
|
|
|
|
|
|
def test_truncate_message_short():
|
|
"""Test that short messages are not truncated."""
|
|
short_message = "short error message"
|
|
result = _truncate_message(short_message)
|
|
assert result == short_message
|
|
|
|
|
|
def test_truncate_message_long():
|
|
"""Test that long messages are truncated."""
|
|
# Create a message longer than the max length
|
|
long_message = "a" * (GRPC_MAX_STATUS_DETAILS_LENGTH + 1000)
|
|
result = _truncate_message(long_message)
|
|
|
|
assert len(result) <= GRPC_MAX_STATUS_DETAILS_LENGTH
|
|
assert result.endswith("... [truncated]")
|
|
|
|
|
|
def test_truncate_message_at_boundary():
|
|
"""Test truncation at the exact boundary."""
|
|
# Create a message exactly at the limit
|
|
exact_message = "a" * GRPC_MAX_STATUS_DETAILS_LENGTH
|
|
result = _truncate_message(exact_message)
|
|
assert result == exact_message
|
|
assert len(result) == GRPC_MAX_STATUS_DETAILS_LENGTH
|
|
|
|
|
|
def test_get_grpc_response_status_truncates_long_message():
|
|
"""Test that long error messages are truncated in INTERNAL errors."""
|
|
long_message = "a" * (GRPC_MAX_STATUS_DETAILS_LENGTH + 1000)
|
|
long_error = RuntimeError(long_message)
|
|
|
|
status = get_grpc_response_status(
|
|
exc=long_error, request_timeout_s=30.0, request_id="test_request_123"
|
|
)
|
|
|
|
assert status.code == grpc.StatusCode.INTERNAL
|
|
assert status.is_error is True
|
|
assert len(status.message) <= GRPC_MAX_STATUS_DETAILS_LENGTH
|
|
assert status.message.endswith("... [truncated]")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|