chore: import upstream snapshot with attribution

python/ray/serve/tests/unit/conftest.py

@@ -0,0 +1,94 @@
+from typing import Tuple
+from unittest.mock import Mock, patch
+
+import pytest
+
+import ray
+from ray._raylet import NodeID
+from ray.serve._private.autoscaling_state import AutoscalingStateManager
+from ray.serve._private.deployment_state import DeploymentStateManager
+from ray.serve._private.test_utils import (
+    MockClusterNodeInfoCache,
+    MockDeploymentActorWrapper,
+    MockKVStore,
+    MockReplicaActorWrapper,
+    MockTimer,
+    dead_replicas_context,
+    replica_rank_context,
+    uninitialized_replicas_context,
+)
+
+
+@pytest.fixture(autouse=True)
+def disallow_ray_init(monkeypatch):
+    def raise_on_init():
+        raise RuntimeError("Unit tests should not depend on Ray being initialized.")
+
+    monkeypatch.setattr(ray, "init", raise_on_init)
+    # Unit tests don't run on a Ray cluster, so stub the runtime context
+    # that callers consult for worker/actor ids.
+    monkeypatch.setattr(ray, "get_runtime_context", Mock())
+
+
+@pytest.fixture
+def mock_deployment_state_manager(
+    request,
+) -> Tuple[DeploymentStateManager, MockTimer, Mock, Mock]:
+    """Fully mocked deployment state manager.
+
+    i.e kv store and gcs client is mocked so we don't need to initialize
+    ray. Also, since this is used for some recovery tests, this yields a
+    method for creating a new mocked deployment state manager.
+    """
+
+    timer = MockTimer()
+    with patch(
+        "ray.serve._private.deployment_state.ActorReplicaWrapper",
+        new=MockReplicaActorWrapper,
+    ), patch(
+        "ray.serve._private.deployment_state.DeploymentActorWrapper",
+        new=MockDeploymentActorWrapper,
+    ), patch(
+        "time.time", new=timer.time
+    ), patch(
+        "ray.serve._private.long_poll.LongPollHost"
+    ) as mock_long_poll, patch(
+        "ray.get_runtime_context"
+    ):
+        kv_store = MockKVStore()
+        cluster_node_info_cache = MockClusterNodeInfoCache()
+        cluster_node_info_cache.add_node(NodeID.from_random().hex())
+        autoscaling_state_manager = AutoscalingStateManager()
+
+        def create_deployment_state_manager(
+            actor_names=None,
+            placement_group_names=None,
+            create_placement_group_fn_override=None,
+        ):
+            if actor_names is None:
+                actor_names = []
+
+            if placement_group_names is None:
+                placement_group_names = []
+
+            return DeploymentStateManager(
+                kv_store,
+                mock_long_poll,
+                actor_names,
+                placement_group_names,
+                cluster_node_info_cache,
+                autoscaling_state_manager,
+                head_node_id_override="fake-head-node-id",
+                create_placement_group_fn_override=create_placement_group_fn_override,
+            )
+
+        yield (
+            create_deployment_state_manager,
+            timer,
+            cluster_node_info_cache,
+            autoscaling_state_manager,
+        )
+
+        dead_replicas_context.clear()
+        replica_rank_context.clear()
+        uninitialized_replicas_context.clear()