fix-lmcache-tests

Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
This commit is contained in:
Sage Ahrac 2025-12-23 13:31:37 +02:00
parent 84818a899e
commit e7d66083ce
No known key found for this signature in database
2 changed files with 11 additions and 0 deletions

View File

@ -25,6 +25,7 @@ def mock_lmcache_engine_event():
lora_id, lora_id,
block_size, block_size,
medium, medium,
lora_name,
): ):
self.block_hashes = block_hashes self.block_hashes = block_hashes
self.parent_block_hash = parent_block_hash self.parent_block_hash = parent_block_hash
@ -32,6 +33,7 @@ def mock_lmcache_engine_event():
self.lora_id = lora_id self.lora_id = lora_id
self.block_size = block_size self.block_size = block_size
self.medium = medium self.medium = medium
self.lora_name = lora_name
return MockEvent( return MockEvent(
block_hashes=["hash1", "hash2"], block_hashes=["hash1", "hash2"],
@ -40,6 +42,7 @@ def mock_lmcache_engine_event():
lora_id=None, lora_id=None,
block_size=16, block_size=16,
medium="GPU", medium="GPU",
lora_name=None,
) )
@ -109,6 +112,7 @@ class TestGetKVConnectorKVCacheEvents:
assert events[0].lora_id is None assert events[0].lora_id is None
assert events[0].block_size == 16 assert events[0].block_size == 16
assert events[0].medium == "GPU" assert events[0].medium == "GPU"
assert events[0].lora_name is None
def test_converts_multiple_events(self, mock_connector): def test_converts_multiple_events(self, mock_connector):
"""Test conversion of multiple events from lmcache engine format.""" """Test conversion of multiple events from lmcache engine format."""
@ -121,6 +125,7 @@ class TestGetKVConnectorKVCacheEvents:
self.lora_id = None self.lora_id = None
self.block_size = 16 self.block_size = 16
self.medium = "GPU" self.medium = "GPU"
self.lora_name = None
events = [MockEvent(i) for i in range(5)] events = [MockEvent(i) for i in range(5)]
mock_connector._lmcache_engine.get_kv_events.return_value = events mock_connector._lmcache_engine.get_kv_events.return_value = events
@ -150,6 +155,7 @@ class TestGetKVConnectorKVCacheEvents:
self.lora_id = 42 self.lora_id = 42
self.block_size = 32 self.block_size = 32
self.medium = "DISK" self.medium = "DISK"
self.lora_name = "lora_example"
mock_connector._lmcache_engine.get_kv_events.return_value = [ mock_connector._lmcache_engine.get_kv_events.return_value = [
MockEventWithLora() MockEventWithLora()
@ -166,6 +172,7 @@ class TestGetKVConnectorKVCacheEvents:
assert event.lora_id == 42 assert event.lora_id == 42
assert event.block_size == 32 assert event.block_size == 32
assert event.medium == "DISK" assert event.medium == "DISK"
assert event.lora_name == "lora_example"
def test_handles_none_parent_block_hash(self, mock_connector): def test_handles_none_parent_block_hash(self, mock_connector):
"""Test handling of events with None parent_block_hash.""" """Test handling of events with None parent_block_hash."""
@ -178,6 +185,7 @@ class TestGetKVConnectorKVCacheEvents:
self.lora_id = None self.lora_id = None
self.block_size = 16 self.block_size = 16
self.medium = "GPU" self.medium = "GPU"
self.lora_name = None
mock_connector._lmcache_engine.get_kv_events.return_value = [ mock_connector._lmcache_engine.get_kv_events.return_value = [
MockEventNoParent() MockEventNoParent()
@ -588,6 +596,7 @@ class TestIntegrationScenarios:
self.lora_id = None self.lora_id = None
self.block_size = 16 self.block_size = 16
self.medium = "GPU" self.medium = "GPU"
self.lora_name = None
# Worker 1 # Worker 1
mock_connector._lmcache_engine.get_kv_events.return_value = [ mock_connector._lmcache_engine.get_kv_events.return_value = [
@ -644,6 +653,7 @@ class TestIntegrationScenarios:
self.lora_id = None self.lora_id = None
self.block_size = 16 self.block_size = 16
self.medium = "GPU" self.medium = "GPU"
self.lora_name = None
for cycle in range(3): for cycle in range(3):
# Get events # Get events

View File

@ -218,6 +218,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
lora_id=e.lora_id, lora_id=e.lora_id,
block_size=e.block_size, block_size=e.block_size,
medium=e.medium, medium=e.medium,
lora_name=e.lora_name,
) )
for e in events for e in events
] ]