wehub-resource-sync
722 lines
25 KiB
Python
722 lines
25 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Unit tests for KVController."""
|
|
|
|
# Standard
|
|
from unittest.mock import Mock, patch
|
|
|
|
# Third Party
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.v1.cache_controller.controllers.kv_controller import KVController
|
|
from lmcache.v1.cache_controller.message import (
|
|
BatchedKVOperationMsg,
|
|
BatchedP2PLookupMsg,
|
|
CheckFinishMsg,
|
|
ClearMsg,
|
|
CompressMsg,
|
|
DecompressMsg,
|
|
KVOpEvent,
|
|
LookupMsg,
|
|
MoveMsg,
|
|
OpType,
|
|
PinMsg,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def kv_controller(mock_reg_controller, mock_cluster_executor):
|
|
"""Create a KVController instance for testing."""
|
|
controller = KVController(registry=mock_reg_controller.registry)
|
|
controller.post_init(
|
|
reg_controller=mock_reg_controller,
|
|
cluster_executor=mock_cluster_executor,
|
|
)
|
|
return controller
|
|
|
|
|
|
class TestKVControllerAdmit:
|
|
"""Test KVController admit operations."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admit_new_key(self, kv_controller):
|
|
"""Test admitting a new KV chunk to the controller."""
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
# Access kv_pool through the registry mock
|
|
report_id = ("test_instance", 0)
|
|
assert report_id in kv_controller.registry.kv_pool
|
|
assert "LocalCPUBackend" in kv_controller.registry.kv_pool[report_id]
|
|
assert 12345 in kv_controller.registry.kv_pool[report_id]["LocalCPUBackend"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admit_duplicate_key(self, kv_controller):
|
|
"""Test admitting same key to different instances."""
|
|
msg1 = BatchedKVOperationMsg(
|
|
instance_id="instance1",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
msg2 = BatchedKVOperationMsg(
|
|
instance_id="instance2",
|
|
worker_id=1,
|
|
location="LocalDiskBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
|
|
await kv_controller.handle_batched_kv_operations(msg1)
|
|
await kv_controller.handle_batched_kv_operations(msg2)
|
|
|
|
report_id1 = ("instance1", 0)
|
|
report_id2 = ("instance2", 1)
|
|
assert 12345 in kv_controller.registry.kv_pool[report_id1]["LocalCPUBackend"]
|
|
assert 12345 in kv_controller.registry.kv_pool[report_id2]["LocalDiskBackend"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admit_multiple_keys(self, kv_controller):
|
|
"""Test admitting multiple different keys."""
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=1000 + i, seq_num=i + 1)
|
|
for i in range(5)
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=operations,
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
report_id = ("test_instance", 0)
|
|
assert report_id in kv_controller.registry.kv_pool
|
|
assert "LocalCPUBackend" in kv_controller.registry.kv_pool[report_id]
|
|
for i in range(5):
|
|
assert (1000 + i) in kv_controller.registry.kv_pool[report_id][
|
|
"LocalCPUBackend"
|
|
]
|
|
|
|
|
|
class TestKVControllerEvict:
|
|
"""Test KVController evict operations."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_evict_existing_key(self, kv_controller):
|
|
"""Test evicting an existing KV chunk."""
|
|
# First admit
|
|
admit_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(admit_msg)
|
|
|
|
# Then evict
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.EVICT, key=12345, seq_num=2)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(evict_msg)
|
|
|
|
report_id = ("test_instance", 0)
|
|
assert (
|
|
report_id not in kv_controller.registry.kv_pool
|
|
or 12345
|
|
not in kv_controller.registry.kv_pool.get(report_id, {}).get(
|
|
"LocalCPUBackend", set()
|
|
)
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_evict_nonexistent_key(self, kv_controller):
|
|
"""Test evicting a non-existent key does nothing."""
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.EVICT, key=99999, seq_num=1)],
|
|
)
|
|
|
|
# Should not raise an error
|
|
await kv_controller.handle_batched_kv_operations(evict_msg)
|
|
# Verify key doesn't exist in registry
|
|
result = kv_controller.registry.find_kv(99999)
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_evict_partial_metadata(self, kv_controller):
|
|
"""Test evicting from one instance while another instance still has it."""
|
|
# Admit two entries for the same key from different instances
|
|
msg1 = BatchedKVOperationMsg(
|
|
instance_id="instance1",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
msg2 = BatchedKVOperationMsg(
|
|
instance_id="instance2",
|
|
worker_id=1,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg1)
|
|
await kv_controller.handle_batched_kv_operations(msg2)
|
|
|
|
# Evict only from the first instance
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="instance1",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.EVICT, key=12345, seq_num=2)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(evict_msg)
|
|
|
|
# Key should still exist in instance2
|
|
report_id1 = ("instance1", 0)
|
|
report_id2 = ("instance2", 1)
|
|
assert (
|
|
report_id1 not in kv_controller.registry.kv_pool
|
|
or 12345
|
|
not in kv_controller.registry.kv_pool[report_id1].get(
|
|
"LocalCPUBackend", set()
|
|
)
|
|
)
|
|
assert 12345 in kv_controller.registry.kv_pool[report_id2]["LocalCPUBackend"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_evict_by_location(self, kv_controller):
|
|
"""Test evicting only matches specific location."""
|
|
# Admit same key to different locations
|
|
msg1 = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
msg2 = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalDiskBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=2)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg1)
|
|
await kv_controller.handle_batched_kv_operations(msg2)
|
|
|
|
# Evict only from CPU backend
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.EVICT, key=12345, seq_num=3)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(evict_msg)
|
|
|
|
# Key should still exist in disk backend
|
|
report_id = ("test_instance", 0)
|
|
assert report_id in kv_controller.registry.kv_pool
|
|
assert "LocalDiskBackend" in kv_controller.registry.kv_pool[report_id]
|
|
assert 12345 in kv_controller.registry.kv_pool[report_id]["LocalDiskBackend"]
|
|
assert (
|
|
"LocalCPUBackend" not in kv_controller.registry.kv_pool[report_id]
|
|
or 12345 not in kv_controller.registry.kv_pool[report_id]["LocalCPUBackend"]
|
|
)
|
|
|
|
|
|
class TestKVControllerSequenceTracking:
|
|
"""Test KVController sequence number tracking."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sequence_number_first_message(self, kv_controller):
|
|
"""Test first message from a source initializes sequence tracker."""
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=5)],
|
|
)
|
|
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
seq_num = kv_controller.reg_controller.registry.get_seq_num(
|
|
"test_instance", 0, "LocalCPUBackend"
|
|
)
|
|
assert seq_num == 5
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sequence_number_continuous(self, kv_controller):
|
|
"""Test continuous sequence numbers don't trigger warnings."""
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=12345 + i, seq_num=i)
|
|
for i in range(1, 6)
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=operations,
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
seq_num = kv_controller.reg_controller.registry.get_seq_num(
|
|
"test_instance", 0, "LocalCPUBackend"
|
|
)
|
|
assert seq_num == 5
|
|
# Sequence discontinuity is now tracked in the registry
|
|
assert kv_controller.registry.get_seq_discontinuity_count() == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sequence_discontinuity_detection(self, kv_controller):
|
|
"""Test detection of gaps in sequence numbers."""
|
|
# First message
|
|
msg1 = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12345, seq_num=1)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg1)
|
|
|
|
# Skip to sequence 5 (gap of 3)
|
|
msg2 = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=12346, seq_num=5)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg2)
|
|
|
|
# Sequence discontinuity is now tracked in the registry
|
|
assert kv_controller.registry.get_seq_discontinuity_count() == 1
|
|
seq_num = kv_controller.reg_controller.registry.get_seq_num(
|
|
"test_instance", 0, "LocalCPUBackend"
|
|
)
|
|
assert seq_num == 5
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sequence_tracking_per_source(self, kv_controller):
|
|
"""Test sequence tracking is independent per source."""
|
|
# Different sources should have independent tracking
|
|
msg1 = BatchedKVOperationMsg(
|
|
instance_id="instance1",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=100, seq_num=1)],
|
|
)
|
|
msg2 = BatchedKVOperationMsg(
|
|
instance_id="instance2",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=200, seq_num=1)],
|
|
)
|
|
|
|
await kv_controller.handle_batched_kv_operations(msg1)
|
|
await kv_controller.handle_batched_kv_operations(msg2)
|
|
|
|
seq_num1 = kv_controller.reg_controller.registry.get_seq_num(
|
|
"instance1", 0, "LocalCPUBackend"
|
|
)
|
|
seq_num2 = kv_controller.reg_controller.registry.get_seq_num(
|
|
"instance2", 0, "LocalCPUBackend"
|
|
)
|
|
|
|
assert seq_num1 == 1
|
|
assert seq_num2 == 1
|
|
|
|
|
|
class TestKVControllerLookup:
|
|
"""Test KVController lookup operations."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lookup_hit(self, kv_controller):
|
|
"""Test successful lookup with prefix match."""
|
|
# Admit some chunks
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=1000 + i, seq_num=i + 1)
|
|
for i in range(3)
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=operations,
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
# Mock token database to return our keys
|
|
with patch.object(
|
|
kv_controller.token_database, "process_tokens"
|
|
) as mock_process:
|
|
mock_process.return_value = [
|
|
(0, 256, 1000),
|
|
(256, 512, 1001),
|
|
(512, 768, 1002),
|
|
]
|
|
|
|
lookup_msg = LookupMsg(
|
|
event_id="event_123",
|
|
tokens=list(range(768)),
|
|
)
|
|
|
|
result = await kv_controller.lookup(lookup_msg)
|
|
|
|
assert result.event_id == "event_123"
|
|
assert "test_instance" in result.layout_info
|
|
assert result.layout_info["test_instance"] == ("LocalCPUBackend", 768)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lookup_miss(self, kv_controller):
|
|
"""Test lookup with no matches."""
|
|
with patch.object(
|
|
kv_controller.token_database, "process_tokens"
|
|
) as mock_process:
|
|
mock_process.return_value = [(0, 256, 9999)]
|
|
|
|
lookup_msg = LookupMsg(
|
|
event_id="event_123",
|
|
tokens=list(range(256)),
|
|
)
|
|
|
|
result = await kv_controller.lookup(lookup_msg)
|
|
|
|
assert result.event_id == "event_123"
|
|
assert len(result.layout_info) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lookup_partial_match(self, kv_controller):
|
|
"""Test lookup with partial prefix match."""
|
|
# Admit only first chunk
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=1000, seq_num=1)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
with patch.object(
|
|
kv_controller.token_database, "process_tokens"
|
|
) as mock_process:
|
|
# Return two chunks but only first exists
|
|
mock_process.return_value = [
|
|
(0, 256, 1000),
|
|
(256, 512, 1001), # This doesn't exist
|
|
]
|
|
|
|
lookup_msg = LookupMsg(
|
|
event_id="event_123",
|
|
tokens=list(range(512)),
|
|
)
|
|
|
|
result = await kv_controller.lookup(lookup_msg)
|
|
|
|
# Should only match up to first chunk
|
|
assert "test_instance" in result.layout_info
|
|
assert result.layout_info["test_instance"] == ("LocalCPUBackend", 256)
|
|
|
|
|
|
class TestKVControllerBatchedP2PLookup:
|
|
"""Test KVController batched P2P lookup operations."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batched_p2p_lookup_success(self, kv_controller):
|
|
"""Test successful batched P2P lookup."""
|
|
# Admit chunks from different instance
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=1000 + i, seq_num=i + 1)
|
|
for i in range(3)
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="remote_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=operations,
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
# Query from different instance
|
|
lookup_msg = BatchedP2PLookupMsg(
|
|
hashes=[1000, 1001, 1002],
|
|
instance_id="query_instance",
|
|
worker_id=0,
|
|
)
|
|
|
|
result = await kv_controller.batched_p2p_lookup(lookup_msg)
|
|
|
|
assert len(result.layout_info) == 1
|
|
instance_id, location, num_hits, peer_url = result.layout_info[0]
|
|
assert instance_id == "remote_instance"
|
|
assert location == "LocalCPUBackend"
|
|
assert num_hits == 3
|
|
assert peer_url == "tcp://localhost:5000"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batched_p2p_lookup_no_match(self, kv_controller):
|
|
"""Test batched P2P lookup with no matches."""
|
|
lookup_msg = BatchedP2PLookupMsg(
|
|
hashes=[9999, 9998, 9997],
|
|
instance_id="query_instance",
|
|
worker_id=0,
|
|
)
|
|
|
|
result = await kv_controller.batched_p2p_lookup(lookup_msg)
|
|
|
|
assert len(result.layout_info) == 1
|
|
instance_id, location, num_hits, peer_url = result.layout_info[0]
|
|
assert instance_id == ""
|
|
assert location == ""
|
|
assert num_hits == 0
|
|
assert peer_url == ""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batched_p2p_lookup_same_instance(self, kv_controller):
|
|
"""Test batched P2P lookup filters out same instance."""
|
|
# Admit chunks from same instance as query
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="query_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=1000, seq_num=1)],
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
lookup_msg = BatchedP2PLookupMsg(
|
|
hashes=[1000],
|
|
instance_id="query_instance",
|
|
worker_id=0,
|
|
)
|
|
|
|
result = await kv_controller.batched_p2p_lookup(lookup_msg)
|
|
|
|
# Should not match same instance
|
|
instance_id, location, num_hits, peer_url = result.layout_info[0]
|
|
assert instance_id == ""
|
|
assert location == ""
|
|
assert num_hits == 0
|
|
assert peer_url == ""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batched_p2p_lookup_partial_match(self, kv_controller):
|
|
"""Test batched P2P lookup with partial matches."""
|
|
# Admit only first two chunks
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=1000 + i, seq_num=i + 1)
|
|
for i in range(2)
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="remote_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=operations,
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
# Query for three chunks
|
|
lookup_msg = BatchedP2PLookupMsg(
|
|
hashes=[1000, 1001, 1002],
|
|
instance_id="query_instance",
|
|
worker_id=0,
|
|
)
|
|
|
|
result = await kv_controller.batched_p2p_lookup(lookup_msg)
|
|
|
|
# Should only match first two
|
|
instance_id, location, num_hits, peer_url = result.layout_info[0]
|
|
assert num_hits == 2
|
|
|
|
|
|
class TestKVControllerDeregister:
|
|
"""Test KVController deregister operations."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deregister_cleanup(self, kv_controller):
|
|
"""Test deregister cleans up KV pool and sequence tracker."""
|
|
# Admit some chunks
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=1000 + i, seq_num=i + 1)
|
|
for i in range(3)
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=operations,
|
|
)
|
|
await kv_controller.handle_batched_kv_operations(msg)
|
|
|
|
# Verify data exists
|
|
report_id = ("test_instance", 0)
|
|
assert report_id in kv_controller.registry.kv_pool
|
|
assert len(kv_controller.registry.kv_pool[report_id]["LocalCPUBackend"]) == 3
|
|
seq_num = kv_controller.reg_controller.registry.get_seq_num(
|
|
"test_instance", 0, "LocalCPUBackend"
|
|
)
|
|
assert seq_num is not None
|
|
|
|
# Deregister directly through registry
|
|
kv_controller.registry.deregister_worker("test_instance", 0)
|
|
|
|
# Verify cleanup
|
|
assert len(kv_controller.registry.kv_pool) == 0
|
|
# Note: seq_tracker is managed by registry, cleaned up by deregister_worker
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deregister_partial_cleanup(self, kv_controller):
|
|
"""Test deregister only removes specific instance-worker."""
|
|
# Admit chunks from different workers
|
|
msg1 = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=0,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=1000, seq_num=1)],
|
|
)
|
|
msg2 = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=1,
|
|
location="LocalCPUBackend",
|
|
operations=[KVOpEvent(op_type=OpType.ADMIT, key=1000, seq_num=1)],
|
|
)
|
|
|
|
await kv_controller.handle_batched_kv_operations(msg1)
|
|
await kv_controller.handle_batched_kv_operations(msg2)
|
|
|
|
# Deregister only worker 0 directly through registry
|
|
kv_controller.registry.deregister_worker("test_instance", 0)
|
|
|
|
# Worker 1's data should remain
|
|
report_id1 = ("test_instance", 1)
|
|
assert report_id1 in kv_controller.registry.kv_pool
|
|
assert 1000 in kv_controller.registry.kv_pool[report_id1]["LocalCPUBackend"]
|
|
|
|
|
|
class TestKVControllerDelegation:
|
|
"""Test KVController delegation to cluster executor."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_clear_delegates_to_executor(self, kv_controller):
|
|
"""Test clear operation delegates to cluster executor."""
|
|
msg = ClearMsg(
|
|
event_id="event_123",
|
|
instance_id="test_instance",
|
|
location="LocalCPUBackend",
|
|
)
|
|
|
|
await kv_controller.clear(msg)
|
|
|
|
kv_controller.cluster_executor.execute.assert_called_once_with("clear", msg)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pin_delegates_to_executor(self, kv_controller):
|
|
"""Test pin operation delegates to cluster executor."""
|
|
msg = PinMsg(
|
|
event_id="event_123",
|
|
instance_id="test_instance",
|
|
location="LocalCPUBackend",
|
|
tokens=[1, 2, 3],
|
|
)
|
|
|
|
await kv_controller.pin(msg)
|
|
|
|
kv_controller.cluster_executor.execute.assert_called_once_with("pin", msg)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compress_delegates_to_executor(self, kv_controller):
|
|
"""Test compress operation delegates to cluster executor."""
|
|
msg = CompressMsg(
|
|
event_id="event_123",
|
|
instance_id="test_instance",
|
|
method="gzip",
|
|
location="LocalCPUBackend",
|
|
tokens=[1, 2, 3],
|
|
)
|
|
|
|
await kv_controller.compress(msg)
|
|
|
|
kv_controller.cluster_executor.execute.assert_called_once_with("compress", msg)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_decompress_delegates_to_executor(self, kv_controller):
|
|
"""Test decompress operation delegates to cluster executor."""
|
|
msg = DecompressMsg(
|
|
event_id="event_123",
|
|
instance_id="test_instance",
|
|
method="gzip",
|
|
location="LocalCPUBackend",
|
|
tokens=[1, 2, 3],
|
|
)
|
|
|
|
await kv_controller.decompress(msg)
|
|
|
|
kv_controller.cluster_executor.execute.assert_called_once_with(
|
|
"decompress", msg
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_move_delegates_to_executor(self, kv_controller):
|
|
"""Test move operation delegates to cluster executor."""
|
|
msg = MoveMsg(
|
|
event_id="event_123",
|
|
old_position=("instance1", "LocalCPUBackend"),
|
|
new_position=("instance2", "LocalCPUBackend"),
|
|
tokens=[1, 2, 3],
|
|
copy=False,
|
|
)
|
|
|
|
await kv_controller.move(msg)
|
|
|
|
kv_controller.cluster_executor.execute.assert_called_once_with("move", msg)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_finish_delegates_to_executor(self, kv_controller):
|
|
"""Test check_finish operation delegates to cluster executor."""
|
|
msg = CheckFinishMsg(event_id="event_123")
|
|
|
|
await kv_controller.check_finish(msg)
|
|
|
|
kv_controller.cluster_executor.execute.assert_called_once_with(
|
|
"check_finish", msg
|
|
)
|
|
|
|
|
|
class TestKVControllerMetrics:
|
|
"""Test KVController metrics setup."""
|
|
|
|
@patch("lmcache.v1.cache_controller.controllers.kv_controller.PrometheusLogger")
|
|
def test_metrics_setup(
|
|
self, mock_prometheus_logger, mock_reg_controller, mock_cluster_executor
|
|
):
|
|
"""Test metrics are properly set up."""
|
|
mock_logger_instance = Mock()
|
|
mock_prometheus_logger.GetInstanceOrNone.return_value = mock_logger_instance
|
|
|
|
_controller = KVController(registry=mock_reg_controller.registry)
|
|
_controller.post_init(
|
|
reg_controller=mock_reg_controller,
|
|
cluster_executor=mock_cluster_executor,
|
|
)
|
|
|
|
# Verify metrics functions were set
|
|
mock_logger_instance.kv_pool_keys_count.set_function.assert_called_once()
|
|
mock_logger_instance.kv_op_seq_discontinuity_count.set_function.assert_called_once()
|
|
|
|
@patch("lmcache.v1.cache_controller.controllers.kv_controller.PrometheusLogger")
|
|
def test_metrics_setup_no_logger(
|
|
self, mock_prometheus_logger, mock_reg_controller, mock_cluster_executor
|
|
):
|
|
"""Test metrics setup when no logger is available."""
|
|
mock_prometheus_logger.GetInstanceOrNone.return_value = None
|
|
|
|
# Should not raise an error
|
|
controller = KVController(registry=mock_reg_controller.registry)
|
|
|
|
assert controller is not None
|