def test_ray_internal_kv_collisions(serve_instance): # noqa: F811
kv1 = RayInternalKVStore()
kv1.put("1", b"1")
assert kv1.get("1") == b"1"
kv2 = RayInternalKVStore("namespace")
assert kv2.get("1") is None
kv2.put("1", b"-1")
assert kv2.get("1") == b"-1"
assert kv1.get("1") == b"1"
def test_ray_internal_kv(serve_instance): # noqa: F811
with pytest.raises(TypeError):
RayInternalKVStore(namespace=1)
RayInternalKVStore(namespace=b"")
kv = RayInternalKVStore()
with pytest.raises(TypeError):
kv.put(1, b"1")
with pytest.raises(TypeError):
kv.put("1", 1)
with pytest.raises(TypeError):
kv.put("1", "1")
kv.put("1", b"2")
assert kv.get("1") == b"2"
kv.put("2", b"4")
assert kv.get("2") == b"4"
kv.put("1", b"3")
assert kv.get("1") == b"3"
assert kv.get("2") == b"4"
def test_ray_internal_kv_timeout(serve_ha): # noqa: F811
# Firstly make sure it's working
kv1 = RayInternalKVStore()
kv1.put("1", b"1")
assert kv1.get("1") == b"1"
# Kill the GCS
ray.worker._global_node.kill_gcs_server()
with pytest.raises(KVStoreError) as e:
kv1.put("2", b"2")
assert e.value.args[0] in (
grpc.StatusCode.UNAVAILABLE,
grpc.StatusCode.DEADLINE_EXCEEDED,
)
