Spaces:
Sleeping
Sleeping
File size: 8,221 Bytes
4b9d59b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | """Tests for RegistryCacheService atomic swap and race condition prevention
This test module verifies:
- Concurrent swap attempts are prevented
- Cache consistency during swap
- Lock timeout behavior
- Error handling
"""
import pytest
import asyncio
from unittest.mock import Mock, AsyncMock, patch
from core.llm.registry.cache import RegistryCacheService, SWAP_LOCK_KEY, LOCK_TTL
@pytest.mark.asyncio
async def test_concurrent_swap_prevention():
"""Test that concurrent swap attempts are prevented by distributed lock"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-concurrent'
# Mock UniversalCacheService
with patch.object(cache.cache, 'get_async', new_callable=AsyncMock) as mock_get, \
patch.object(cache.cache, 'set_async', new_callable=AsyncMock) as mock_set, \
patch.object(cache.cache, 'delete_async', new_callable=AsyncMock):
# First call - no lock exists
mock_get.return_value = None # No existing lock
mock_set.return_value = True
models = [
{'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 8192}
]
# First swap should succeed
result1 = await cache.atomic_swap_registry(tenant_id, models)
assert result1 is True
# Verify lock was acquired (check first call to set_async)
assert mock_set.call_count > 0
first_call = mock_set.call_args_list[0]
assert first_call[0][0] == SWAP_LOCK_KEY # First argument is lock key
assert first_call[0][1] == "swapping"
@pytest.mark.asyncio
async def test_concurrent_swap_blocks_second_attempt():
"""Test that second swap attempt fails when lock is held"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-block'
# Mock UniversalCacheService
with patch.object(cache.cache, 'get_async', new_callable=AsyncMock) as mock_get:
# Lock already exists
mock_get.return_value = "swapping"
models = [
{'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 8192}
]
# Second swap should raise exception
with pytest.raises(Exception) as exc_info:
await cache.atomic_swap_registry(tenant_id, models)
assert "Swap in progress" in str(exc_info.value)
@pytest.mark.asyncio
async def test_cache_consistency_during_swap():
"""Test that cache remains consistent during atomic swap"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-consistency'
# Pre-populate cache with old data
old_models = [
{'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 8192}
]
new_models = [
{'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 128000}, # Updated
{'provider': 'openai', 'model_name': 'gpt-4-turbo', 'context_window': 128000} # New
]
with patch.object(cache.cache, 'get_async', new_callable=AsyncMock, return_value=None), \
patch.object(cache.cache, 'set_async', new_callable=AsyncMock, return_value=True), \
patch.object(cache.cache, 'delete_async', new_callable=AsyncMock) as mock_delete:
# Perform atomic swap
result = await cache.atomic_swap_registry(tenant_id, new_models)
assert result is True
# Verify all new models were cached (2 models + 2 lists = 4 calls)
assert cache.cache.set_async.call_count >= len(new_models)
# Verify lock was released
mock_delete.assert_called_once()
delete_call_args = mock_delete.call_args[0]
assert delete_call_args[0] == SWAP_LOCK_KEY
assert delete_call_args[1] == tenant_id
@pytest.mark.asyncio
async def test_lock_timeout_behavior():
"""Test that lock expires after timeout"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-timeout'
# Mock cache that simulates lock timeout
lock_acquired = False
async def mock_get(key, tenant):
if key == SWAP_LOCK_KEY and lock_acquired:
# Simulate lock timeout on second check
return None
return "swapping" if lock_acquired else None
async def mock_set(key, value, ttl, tenant):
nonlocal lock_acquired
if key == SWAP_LOCK_KEY:
lock_acquired = True
return True
with patch.object(cache.cache, 'get_async', side_effect=mock_get), \
patch.object(cache.cache, 'set_async', side_effect=mock_set), \
patch.object(cache.cache, 'delete_async', new_callable=AsyncMock):
models = [{'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 8192}]
# First swap should acquire lock
result1 = await cache.atomic_swap_registry(tenant_id, models)
assert result1 is True
@pytest.mark.asyncio
async def test_swap_error_handling():
"""Test error handling during swap"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-error'
# Mock cache that fails during set (after lock is acquired)
call_count = [0]
async def mock_set_with_error(key, value, ttl, tenant):
call_count[0] += 1
if key != SWAP_LOCK_KEY: # Fail on model set, not lock set
raise Exception("Redis error")
return True
with patch.object(cache.cache, 'get_async', new_callable=AsyncMock, return_value=None), \
patch.object(cache.cache, 'set_async', side_effect=mock_set_with_error), \
patch.object(cache.cache, 'delete_async', new_callable=AsyncMock) as mock_delete:
models = [{'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 8192}]
# Swap should handle error gracefully (catches and logs exceptions)
result = await cache.atomic_swap_registry(tenant_id, models)
# Result should still be True (errors are logged, not raised)
assert result is True
# Verify lock was still released despite errors
mock_delete.assert_called_once()
delete_call_args = mock_delete.call_args[0]
assert delete_call_args[0] == SWAP_LOCK_KEY
assert delete_call_args[1] == tenant_id
@pytest.mark.asyncio
async def test_warm_cache_no_lock_required():
"""Test that warm_cache doesn't require lock"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-warm'
with patch.object(cache.cache, 'set_async', new_callable=AsyncMock, return_value=True):
models = [
{'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 8192},
{'provider': 'anthropic', 'model_name': 'claude-3-opus', 'context_window': 200000}
]
# warm_cache should not use lock
result = await cache.warm_cache(tenant_id, models)
# Should complete without error
assert result is None # warm_cache returns None
@pytest.mark.asyncio
async def test_invalidate_tenant_clears_all_keys():
"""Test that invalidate_tenant clears all tenant cache keys"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-invalidate'
with patch.object(cache.cache, 'delete_tenant_all', new_callable=AsyncMock, return_value=42):
count = await cache.invalidate_tenant(tenant_id)
assert count == 42
@pytest.mark.asyncio
async def test_model_cache_operations():
"""Test basic model cache operations"""
cache = RegistryCacheService()
tenant_id = 'test-tenant-model-ops'
with patch.object(cache.cache, 'get_async', new_callable=AsyncMock) as mock_get, \
patch.object(cache.cache, 'set_async', new_callable=AsyncMock, return_value=True) as mock_set:
# Set model
model_data = {'provider': 'openai', 'model_name': 'gpt-4', 'context_window': 8192}
await cache.set_model(tenant_id, 'openai', 'gpt-4', model_data)
# Verify set was called
mock_set.assert_called_once()
# Get model (cache hit)
mock_get.return_value = model_data
result = await cache.get_model(tenant_id, 'openai', 'gpt-4')
assert result == model_data
# Get model (cache miss)
mock_get.return_value = None
result = await cache.get_model(tenant_id, 'openai', 'gpt-5')
assert result is None
if __name__ == '__main__':
pytest.main([__file__, '-v'])
|