chore: import upstream snapshot with attribution

python/ray/serve/tests/test_proxy_actor_wrapper.py

@@ -0,0 +1,270 @@
+import asyncio
+import concurrent.futures
+from unittest.mock import Mock
+
+import pytest
+
+from ray.exceptions import RayTaskError
+from ray.serve._private.proxy_state import ActorProxyWrapper, wrap_as_future
+from ray.serve.schema import LoggingConfig
+
+
+def _create_object_ref_mock():
+    fut = concurrent.futures.Future()
+
+    ray_object_ref_mock = Mock(future=Mock(return_value=fut))
+
+    return ray_object_ref_mock, fut
+
+
+def _create_mocked_actor_proxy_wrapper(actor_handle_mock):
+    return ActorProxyWrapper(
+        logging_config=LoggingConfig(),
+        actor_handle=actor_handle_mock,
+        node_id="some_node_id",
+    )
+
+
+@pytest.mark.asyncio
+async def test_wrap_as_future_timeout():
+    object_ref_mock, fut = _create_object_ref_mock()
+
+    aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=0)
+
+    assert not aio_fut.done()
+    # Yield the event-loop
+    await asyncio.sleep(0.001)
+
+    assert aio_fut.done()
+    with pytest.raises(TimeoutError) as exc_info:
+        aio_fut.result()
+
+    assert "Future cancelled after timeout 0s" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_wrap_as_future_success():
+    # Test #1: Validate wrapped asyncio future completes, upon completion of the
+    #          ObjectRef's one
+
+    object_ref_mock, fut = _create_object_ref_mock()
+
+    aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=3600)
+
+    assert not aio_fut.done()
+    # Complete source future (ObjectRef one)
+    fut.set_result("test")
+    # Yield the event-loop
+    await asyncio.sleep(0.001)
+
+    assert aio_fut.done()
+    assert aio_fut.result() == "test"
+
+    # Test #2: Wrapped asyncio future completed before time-out expiration,
+    #          should not be affected by the cancellation callback
+
+    object_ref_mock, fut = _create_object_ref_mock()
+    # Purposefully set timeout to 0, ie future has to be cancelled upon next
+    # event-loop iteration
+    aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=0)
+
+    assert not aio_fut.done()
+    # Complete source future (ObjectRef one)
+    fut.set_result("test")
+    # Yield the event-loop
+    await asyncio.sleep(0.001)
+
+    assert aio_fut.done()
+    assert aio_fut.result() == "test"
+
+
+@pytest.mark.asyncio
+async def test_wrap_as_future_timeout_does_not_assert_on_late_source_completion(
+    monkeypatch,
+):
+    object_ref_mock, fut = _create_object_ref_mock()
+    loop = asyncio.get_running_loop()
+    queued_callbacks = []
+
+    def queue_callback(callback, *args, context=None):
+        queued_callbacks.append((callback, args, context))
+
+    monkeypatch.setattr(loop, "call_soon_threadsafe", queue_callback)
+
+    aio_fut = wrap_as_future(ref=object_ref_mock, timeout_s=0)
+
+    fut.set_result("test")
+    assert len(queued_callbacks) == 1
+
+    await asyncio.sleep(0.001)
+
+    assert aio_fut.done()
+    with pytest.raises(TimeoutError):
+        aio_fut.result()
+
+    callback, args, _ = queued_callbacks.pop()
+    callback(*args)
+
+
+@pytest.mark.parametrize(
+    ("response", "is_ready"),
+    [
+        # ProxyActor.ready responds with an tuple/array of 2 strings
+        ('["foo", "bar"]', True),
+        ("malformed_json", False),
+        (Exception(), False),
+    ],
+)
+@pytest.mark.asyncio
+async def test_is_ready_check_success(response, is_ready):
+    """Tests calling is_ready method on ProxyActorWrapper, mocking out underlying
+    ActorHandle response"""
+    object_ref_mock, fut = _create_object_ref_mock()
+    actor_handle_mock = Mock(ready=Mock(remote=Mock(return_value=object_ref_mock)))
+
+    proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
+
+    for _ in range(10):
+        assert proxy_wrapper.is_ready(timeout_s=1) is None
+        # Yield loop!
+        await asyncio.sleep(0.01)
+
+    # Complete source future
+    if isinstance(response, Exception):
+        object_ref_mock.future().set_exception(response)
+    else:
+        object_ref_mock.future().set_result(response)
+
+    # Yield loop!
+    await asyncio.sleep(0)
+    # NOTE: Timeout setting is only relevant, in case there's no pending request
+    #       and one will be issued
+    assert proxy_wrapper.is_ready(timeout_s=1) is is_ready
+
+
+@pytest.mark.asyncio
+async def test_is_ready_check_timeout():
+    object_ref_mock, fut = _create_object_ref_mock()
+    actor_handle_mock = Mock(ready=Mock(remote=Mock(return_value=object_ref_mock)))
+
+    proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
+
+    # First call, invokes ProxyActor.ready call
+    assert proxy_wrapper.is_ready(timeout_s=0) is None
+    # Yield loop!
+    await asyncio.sleep(0.001)
+
+    assert proxy_wrapper.is_ready(timeout_s=0) is False
+
+
+@pytest.mark.parametrize(
+    ("response", "is_healthy"),
+    [
+        (True, True),
+        (RayTaskError("check_health", "<traceback>", "cuz"), False),
+    ],
+)
+@pytest.mark.asyncio
+async def test_is_healthy_check_success(response, is_healthy):
+    """Tests calling is_healthy method on ProxyActorWrapper, mocking out underlying
+    ActorHandle response"""
+    object_ref_mock, fut = _create_object_ref_mock()
+    actor_handle_mock = Mock(
+        check_health=Mock(remote=Mock(return_value=object_ref_mock))
+    )
+
+    proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
+
+    for _ in range(10):
+        assert proxy_wrapper.is_healthy(timeout_s=1) is None
+        # Yield loop!
+        await asyncio.sleep(0.01)
+
+    object_ref_mock.future.assert_called_once()
+
+    # Complete source future
+    if isinstance(response, Exception):
+        object_ref_mock.future.return_value.set_exception(response)
+    else:
+        object_ref_mock.future.return_value.set_result(response)
+
+    # Yield loop!
+    await asyncio.sleep(0)
+    # NOTE: Timeout setting is only relevant, in case there's no pending request
+    #       and one will be issued
+    assert proxy_wrapper.is_healthy(timeout_s=1) is is_healthy
+
+
+@pytest.mark.asyncio
+async def test_is_healthy_check_timeout():
+    object_ref_mock, fut = _create_object_ref_mock()
+    actor_handle_mock = Mock(
+        check_health=Mock(remote=Mock(return_value=object_ref_mock))
+    )
+
+    proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
+
+    # First call, invokes ProxyActor.ready call
+    assert proxy_wrapper.is_healthy(timeout_s=0) is None
+    # Yield loop!
+    await asyncio.sleep(0.001)
+
+    assert proxy_wrapper.is_healthy(timeout_s=0) is False
+
+
+@pytest.mark.parametrize(
+    ("response", "is_drained"),
+    [
+        (True, True),
+        (False, False),
+        (RayTaskError("is_drained", "<traceback>", "cuz"), False),
+    ],
+)
+@pytest.mark.asyncio
+async def test_is_drained_check_success(response, is_drained):
+    """Tests calling is_drained method on ProxyActorWrapper, mocking out underlying
+    ActorHandle response"""
+    object_ref_mock, fut = _create_object_ref_mock()
+    actor_handle_mock = Mock(is_drained=Mock(remote=Mock(return_value=object_ref_mock)))
+
+    proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
+
+    for _ in range(10):
+        assert proxy_wrapper.is_drained(timeout_s=1) is None
+        # Yield loop!
+        await asyncio.sleep(0.01)
+
+    object_ref_mock.future.assert_called_once()
+
+    # Complete source future
+    if isinstance(response, Exception):
+        object_ref_mock.future.return_value.set_exception(response)
+    else:
+        object_ref_mock.future.return_value.set_result(response)
+
+    # Yield loop!
+    await asyncio.sleep(0)
+    # NOTE: Timeout setting is only relevant, in case there's no pending request
+    #       and one will be issued
+    assert proxy_wrapper.is_drained(timeout_s=1) is is_drained
+
+
+@pytest.mark.asyncio
+async def test_is_drained_check_timeout():
+    object_ref_mock, fut = _create_object_ref_mock()
+    actor_handle_mock = Mock(is_drained=Mock(remote=Mock(return_value=object_ref_mock)))
+
+    proxy_wrapper = _create_mocked_actor_proxy_wrapper(actor_handle_mock)
+
+    # First call, invokes ProxyActor.ready call
+    assert proxy_wrapper.is_drained(timeout_s=0) is None
+    # Yield loop!
+    await asyncio.sleep(0.001)
+
+    assert proxy_wrapper.is_drained(timeout_s=0) is False
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-s", __file__]))