add more test

Signed-off-by: Gene Su <e870252314@gmail.com>

python/ray/serve/BUILD

@@ -665,3 +665,11 @@ py_test(
     tags = ["exclusive", "team:serve"],
     deps = [":serve_lib"],
 )
+
+py_test(
+    name = "test_endpoint_state",
+    size = "small",
+    srcs = serve_tests_srcs,
+    tags = ["exclusive", "team:serve"],
+    deps = [":serve_lib"],
+)

python/ray/serve/tests/test_application_state.py

@@ -659,13 +659,12 @@ def test_recover_during_update(mocked_application_state_manager):
 
 
 def test_is_shutdown(mocked_application_state_manager):
-    """Test is_shutdown() returns True the correct state.
-
+    """Test `is_shutdown()` returns the correct state.
 
     When shutting down applications before deployments are deleted, application state
-    `is_deleted()` should return false and `is_shutdown()` should return false.
+    `is_deleted()` should return False and `is_shutdown()` should return False.
     When shutting down applications after deployments are deleted, application state
-    `is_deleted()` should return true and `is_shutdown()` should return true
+    `is_deleted()` should return True and `is_shutdown()` should return True
     """
     (
         app_state_manager,
@@ -689,13 +688,13 @@ def test_is_shutdown(mocked_application_state_manager):
     assert app_state.status == ApplicationStatus.RUNNING
 
     # When shutting down applications before deployments are deleted, application state
-    # `is_deleted()` should return false and `is_shutdown()` should return false
+    # `is_deleted()` should return False and `is_shutdown()` should return False
     app_state_manager.shutdown()
     assert not app_state.is_deleted()
     assert not app_state_manager.is_shutdown()
 
     # When shutting down applications after deployments are deleted, application state
-    # `is_deleted()` should return true and `is_shutdown()` should return true
+    # `is_deleted()` should return True and `is_shutdown()` should return True
     deployment_state_manager.delete_deployment(deployment_name)
     deployment_state_manager.set_deployment_deleted(deployment_name)
     app_state_manager.update()

python/ray/serve/tests/test_deployment_state.py

@@ -2537,9 +2537,6 @@ def test_shutdown(mock_deployment_state_manager, is_driver_deployment):
 
     deployment_state_manager.shutdown()
 
-    # After shutdown, `is_shutdown()` should return True
-    assert not deployment_state_manager.is_shutdown()
-
     timer.advance(grace_period_s + 0.1)
     deployment_state_manager.update()
 
@@ -2554,6 +2551,9 @@ def test_shutdown(mock_deployment_state_manager, is_driver_deployment):
     check_counts(deployment_state, total=0)
     assert len(deployment_state_manager.get_deployment_statuses()) == 0
 
+    # After all deployments shutdown, `is_shutdown()` should return True
+    assert deployment_state_manager.is_shutdown()
+
 
 def test_stopping_replicas_ranking():
     @dataclass

python/ray/serve/tests/test_endpoint_state.py

@@ -0,0 +1,65 @@
+import sys
+from typing import Any, Tuple
+from unittest.mock import patch, Mock
+
+import pytest
+from ray.serve._private.endpoint_state import EndpointState
+
+
+class MockKVStore:
+    def __init__(self):
+        self.store = dict()
+
+    def put(self, key: str, val: Any) -> bool:
+        if not isinstance(key, str):
+            raise TypeError("key must be a string, got: {}.".format(type(key)))
+        self.store[key] = val
+        return True
+
+    def get(self, key: str) -> Any:
+        if not isinstance(key, str):
+            raise TypeError("key must be a string, got: {}.".format(type(key)))
+        return self.store.get(key, None)
+
+    def delete(self, key: str) -> bool:
+        if not isinstance(key, str):
+            raise TypeError("key must be a string, got: {}.".format(type(key)))
+
+        if key in self.store:
+            del self.store[key]
+            return True
+
+        return False
+
+
+@pytest.fixture
+def mock_endpoint_state() -> Tuple[EndpointState, Mock]:
+    with patch("ray.serve._private.long_poll.LongPollHost") as mock_long_poll:
+        endpoint_state = EndpointState(
+            kv_store=MockKVStore(),
+            long_poll_host=mock_long_poll,
+        )
+        yield endpoint_state
+
+
+def test_is_shutdown(mock_endpoint_state):
+    """Test `is_shutdown()` returns the correct state.
+
+    Before shutting down endpoint `is_shutdown()` should return False.
+    After shutting down endpoint `is_shutdown()` should return True.
+    """
+    # Setup endpoint state with checkpoint
+    endpoint_state = mock_endpoint_state
+    endpoint_state._checkpoint()
+
+    # Before shutdown is called, `is_shutdown()` should return False
+    assert not endpoint_state.is_shutdown()
+
+    endpoint_state.shutdown()
+
+    # After shutdown is called, `is_shutdown()` should return True
+    assert endpoint_state.is_shutdown()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

python/ray/serve/tests/test_standalone3.py

@@ -509,5 +509,33 @@ def _check():
     serve.shutdown()
 
 
+# def test_controller_shutdown_gracefully(
+#     shutdown_ray, call_ray_stop_only  # noqa: F811
+# ):
+#     """TODO
+#     """
+#     # Setup a cluster with 2 nodes
+#     cluster = Cluster()
+#     cluster.add_node()
+#     cluster.add_node()
+#     cluster.wait_for_nodes()
+#     ray.init(address=cluster.address)
+#
+#     # Deploy 2 replicas
+#     @serve.deployment(num_replicas=2)
+#     class HelloModel:
+#         def __call__(self):
+#             return "hello"
+#
+#     model = HelloModel.bind()
+#     serve.run(target=model)
+#
+#     # Ensure total actors of 2 proxies, 1 controller, and 2 replicas
+#     wait_for_condition(lambda: len(ray._private.state.actors()) == 5)
+#     assert len(ray.nodes()) == 2
+#
+#     serve.shutdown()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-s", __file__]))