120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Unit tests for store policy interface and DefaultStorePolicy.
|
|
|
|
Tests are written against the StorePolicy contract defined in store_policy.py.
|
|
"""
|
|
|
|
# Third Party
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import ObjectKey
|
|
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig
|
|
from lmcache.v1.distributed.storage_controllers.store_policy import (
|
|
AdapterDescriptor,
|
|
DefaultStorePolicy,
|
|
)
|
|
|
|
# =============================================================================
|
|
# Helpers
|
|
# =============================================================================
|
|
|
|
|
|
def make_object_key(chunk_id: int) -> ObjectKey:
|
|
"""Create a test ObjectKey with the given chunk ID."""
|
|
return ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
|
|
model_name="test_model",
|
|
kv_rank=0,
|
|
)
|
|
|
|
|
|
def make_descriptor(index: int) -> AdapterDescriptor:
|
|
"""Create an AdapterDescriptor for testing."""
|
|
config = MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=10.0)
|
|
return AdapterDescriptor(index=index, config=config)
|
|
|
|
|
|
# =============================================================================
|
|
# DefaultStorePolicy Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestDefaultStorePolicyTargets:
|
|
"""Test DefaultStorePolicy.select_store_targets behavior."""
|
|
|
|
def test_single_adapter_all_keys(self):
|
|
"""All keys should be sent to the single adapter."""
|
|
policy = DefaultStorePolicy()
|
|
keys = [make_object_key(i) for i in range(3)]
|
|
adapters = [make_descriptor(0)]
|
|
|
|
result = policy.select_store_targets(keys, adapters)
|
|
|
|
assert 0 in result
|
|
assert result[0] == keys
|
|
|
|
def test_multiple_adapters_all_keys_to_each(self):
|
|
"""All keys should be sent to every adapter."""
|
|
policy = DefaultStorePolicy()
|
|
keys = [make_object_key(i) for i in range(3)]
|
|
adapters = [make_descriptor(0), make_descriptor(1)]
|
|
|
|
result = policy.select_store_targets(keys, adapters)
|
|
|
|
assert len(result) == 2
|
|
assert result[0] == keys
|
|
assert result[1] == keys
|
|
|
|
def test_empty_adapters_returns_empty(self):
|
|
"""No adapters means no store targets."""
|
|
policy = DefaultStorePolicy()
|
|
keys = [make_object_key(0)]
|
|
|
|
result = policy.select_store_targets(keys, [])
|
|
|
|
assert result == {}
|
|
|
|
def test_empty_keys_returns_empty_lists(self):
|
|
"""Empty keys list should produce empty lists for each adapter."""
|
|
policy = DefaultStorePolicy()
|
|
adapters = [make_descriptor(0)]
|
|
|
|
result = policy.select_store_targets([], adapters)
|
|
|
|
assert 0 in result
|
|
assert result[0] == []
|
|
|
|
def test_returns_copies_not_references(self):
|
|
"""Returned lists should be independent copies of the input."""
|
|
policy = DefaultStorePolicy()
|
|
keys = [make_object_key(0)]
|
|
adapters = [make_descriptor(0)]
|
|
|
|
result = policy.select_store_targets(keys, adapters)
|
|
|
|
# Mutating the result should not affect the input
|
|
result[0].append(make_object_key(99))
|
|
assert len(keys) == 1
|
|
|
|
|
|
class TestDefaultStorePolicyDeletions:
|
|
"""Test DefaultStorePolicy.select_l1_deletions behavior."""
|
|
|
|
def test_never_deletes(self):
|
|
"""DefaultStorePolicy should never delete from L1."""
|
|
policy = DefaultStorePolicy()
|
|
keys = [make_object_key(i) for i in range(5)]
|
|
|
|
result = policy.select_l1_deletions(keys)
|
|
|
|
assert result == []
|
|
|
|
def test_empty_keys_returns_empty(self):
|
|
"""Empty input should return empty output."""
|
|
policy = DefaultStorePolicy()
|
|
|
|
result = policy.select_l1_deletions([])
|
|
|
|
assert result == []
|