Spaces:
Paused
Paused
File size: 2,774 Bytes
a5784e9 | 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 | import os
import sys
import unittest
from datetime import datetime, timedelta
from unittest.mock import patch
# Add repository root to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
# Import AFTER sys.path modification
from api_utils.utils_ext import cooldown_manager
class TestCooldownPersistence(unittest.TestCase):
def setUp(self):
# Create a temporary file for testing
self.test_file = 'test_cooldown_status.json'
self.patcher = patch('api_utils.utils_ext.cooldown_manager.COOLDOWN_FILE', self.test_file)
self.mock_file = self.patcher.start()
def tearDown(self):
self.patcher.stop()
if os.path.exists(self.test_file):
os.remove(self.test_file)
def test_legacy_persistence(self):
"""Test saving and loading flat structure (legacy behavior)"""
now = datetime.now()
profiles = {
'profile1': now,
'profile2': now - timedelta(hours=1)
}
cooldown_manager.save_cooldown_profiles(profiles)
loaded_profiles = cooldown_manager.load_cooldown_profiles()
self.assertIn('profile1', loaded_profiles)
self.assertIsInstance(loaded_profiles['profile1'], datetime)
# Compare ISO strings to avoid microsecond precision issues during serialization
self.assertEqual(loaded_profiles['profile1'].isoformat(), now.isoformat())
def test_nested_persistence(self):
"""Test saving and loading nested structure (model-specific cooldowns)"""
now = datetime.now()
profiles = {
'profile1': {
'gpt-4': now,
'claude-3': now - timedelta(hours=1)
},
'profile2': now
}
cooldown_manager.save_cooldown_profiles(profiles)
loaded_profiles = cooldown_manager.load_cooldown_profiles()
# Check nested structure
self.assertIn('profile1', loaded_profiles)
self.assertIsInstance(loaded_profiles['profile1'], dict)
self.assertIn('gpt-4', loaded_profiles['profile1'])
self.assertEqual(loaded_profiles['profile1']['gpt-4'].isoformat(), now.isoformat())
# Check mixed flat structure
self.assertIn('profile2', loaded_profiles)
self.assertIsInstance(loaded_profiles['profile2'], datetime)
self.assertEqual(loaded_profiles['profile2'].isoformat(), now.isoformat())
def test_corrupted_data_handling(self):
"""Test handling of corrupted data"""
with open(self.test_file, 'w') as f:
f.write("{invalid_json")
loaded_profiles = cooldown_manager.load_cooldown_profiles()
self.assertEqual(loaded_profiles, {})
if __name__ == '__main__':
unittest.main()
|