| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from unittest.mock import AsyncMock, MagicMock, patch |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | grpc = pytest.importorskip("grpc") |
| 8 | health_pb2 = pytest.importorskip("grpc_health.v1.health_pb2") |
| 9 | VllmHealthServicer = pytest.importorskip( |
| 10 | "smg_grpc_servicer.vllm.health_servicer" |
| 11 | ).VllmHealthServicer |
| 12 | |
| 13 | SERVING = health_pb2.HealthCheckResponse.SERVING |
| 14 | NOT_SERVING = health_pb2.HealthCheckResponse.NOT_SERVING |
| 15 | SERVICE_UNKNOWN = health_pb2.HealthCheckResponse.SERVICE_UNKNOWN |
| 16 | |
| 17 | |
| 18 | @pytest.fixture |
| 19 | def async_llm(): |
| 20 | mock = MagicMock() |
| 21 | mock.check_health = AsyncMock() |
| 22 | return mock |
| 23 | |
| 24 | |
| 25 | @pytest.fixture |
| 26 | def context(): |
| 27 | return MagicMock(spec=grpc.aio.ServicerContext) |
| 28 | |
| 29 | |
| 30 | @pytest.fixture |
| 31 | def servicer(async_llm): |
| 32 | return VllmHealthServicer(async_llm) |
| 33 | |
| 34 | |
| 35 | @pytest.fixture |
| 36 | def request_msg(): |
| 37 | msg = MagicMock() |
| 38 | msg.service = "" |
| 39 | return msg |
| 40 | |
| 41 | |
| 42 | # -- Check() tests -- |
| 43 | |
| 44 | |
| 45 | @pytest.mark.asyncio |
| 46 | async def test_check_serving_overall(servicer, request_msg, context, async_llm): |
| 47 | request_msg.service = "" |
| 48 | response = await servicer.Check(request_msg, context) |
| 49 | assert response.status == SERVING |
| 50 | async_llm.check_health.assert_awaited_once() |
| 51 | |
| 52 | |
| 53 | @pytest.mark.asyncio |
| 54 | async def test_check_serving_vllm_service(servicer, request_msg, context, async_llm): |
| 55 | request_msg.service = "vllm.grpc.engine.VllmEngine" |
| 56 | response = await servicer.Check(request_msg, context) |
| 57 | assert response.status == SERVING |
| 58 | async_llm.check_health.assert_awaited_once() |
| 59 | |
| 60 | |
| 61 | @pytest.mark.asyncio |
| 62 | async def test_check_not_serving_engine_errored( |
| 63 | servicer, request_msg, context, async_llm |
| 64 | ): |
| 65 | async_llm.check_health = AsyncMock(side_effect=Exception("engine dead")) |
| 66 | request_msg.service = "" |
| 67 | response = await servicer.Check(request_msg, context) |
| 68 | assert response.status == NOT_SERVING |
| 69 | |
| 70 | |
| 71 | @pytest.mark.asyncio |
| 72 | async def test_check_not_serving_shutting_down( |
| 73 | servicer, request_msg, context, async_llm |
| 74 | ): |
| 75 | servicer.set_not_serving() |
| 76 | request_msg.service = "" |
| 77 | response = await servicer.Check(request_msg, context) |
| 78 | assert response.status == NOT_SERVING |
| 79 | async_llm.check_health.assert_not_awaited() |
| 80 | |
| 81 | |
| 82 | @pytest.mark.asyncio |
| 83 | async def test_check_unknown_service_status(servicer, request_msg, context): |
| 84 | request_msg.service = "nonexistent.Service" |
| 85 | response = await servicer.Check(request_msg, context) |
| 86 | assert response.status == SERVICE_UNKNOWN |
| 87 | |
| 88 | |
| 89 | @pytest.mark.asyncio |
| 90 | async def test_check_unknown_service_grpc_code(servicer, request_msg, context): |
| 91 | request_msg.service = "fake.Svc" |
| 92 | await servicer.Check(request_msg, context) |
| 93 | context.set_code.assert_called_once_with(grpc.StatusCode.NOT_FOUND) |
| 94 | context.set_details.assert_called_once() |
| 95 | details_arg = context.set_details.call_args[0][0] |
| 96 | assert "fake.Svc" in details_arg |
| 97 | |
| 98 | |
| 99 | @pytest.mark.asyncio |
| 100 | @patch("smg_grpc_servicer.vllm.health_servicer.logger") |
| 101 | async def test_check_logs_exception_on_error( |
| 102 | mock_logger, servicer, request_msg, context, async_llm |
| 103 | ): |
| 104 | async_llm.check_health = AsyncMock(side_effect=Exception("engine exploded")) |
| 105 | request_msg.service = "" |
| 106 | await servicer.Check(request_msg, context) |
| 107 | mock_logger.exception.assert_called_once() |
| 108 | log_args = mock_logger.exception.call_args |
| 109 | assert "service" in str(log_args).lower() |
| 110 | |
| 111 | |
| 112 | # -- Watch() tests -- |
| 113 | |
| 114 | |
| 115 | @pytest.mark.asyncio |
| 116 | async def test_watch_yields_serving(servicer, request_msg, context, async_llm): |
| 117 | request_msg.service = "" |
| 118 | watch_iter = servicer.Watch(request_msg, context) |
| 119 | first = await anext(watch_iter.__aiter__()) |
| 120 | assert first.status == SERVING |
| 121 | |
| 122 | |
| 123 | @pytest.mark.asyncio |
| 124 | async def test_watch_yields_not_serving(servicer, request_msg, context, async_llm): |
| 125 | async_llm.check_health = AsyncMock(side_effect=Exception("engine down")) |
| 126 | request_msg.service = "" |
| 127 | watch_iter = servicer.Watch(request_msg, context) |
| 128 | first = await anext(watch_iter.__aiter__()) |
| 129 | assert first.status == NOT_SERVING |
| 130 | |
| 131 | |
| 132 | @pytest.mark.asyncio |
| 133 | async def test_watch_unknown_service(servicer, request_msg, context): |
| 134 | request_msg.service = "fake.Service" |
| 135 | results = [] |
| 136 | async for response in servicer.Watch(request_msg, context): |
| 137 | results.append(response) |
| 138 | assert len(results) == 1 |
| 139 | assert results[0].status == SERVICE_UNKNOWN |
| 140 | # Watch returns SERVICE_UNKNOWN in the response body (not as a gRPC error |
| 141 | # code) so the stream terminates normally -- unlike Check, which sets |
| 142 | # NOT_FOUND on the gRPC context for unknown services. |
| 143 | context.set_code.assert_not_called() |