Spaces:
Paused
Paused
| import asyncio | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from unittest.mock import MagicMock, patch | |
| import httpx | |
| import pytest | |
| import respx | |
| from fastapi.testclient import TestClient | |
| sys.path.insert( | |
| 0, os.path.abspath("../../..") | |
| ) # Adds the parent directory to the system path | |
| from unittest.mock import AsyncMock | |
| from litellm.caching.in_memory_cache import InMemoryCache | |
| def test_in_memory_openai_obj_cache(): | |
| from openai import OpenAI | |
| openai_obj = OpenAI(api_key="my-fake-key") | |
| in_memory_cache = InMemoryCache() | |
| in_memory_cache.set_cache(key="my-fake-key", value=openai_obj) | |
| cached_obj = in_memory_cache.get_cache(key="my-fake-key") | |
| assert cached_obj is not None | |
| assert cached_obj == openai_obj | |
| def test_in_memory_cache_max_size_per_item(): | |
| """ | |
| Test that the cache will not store items larger than the max size per item | |
| """ | |
| in_memory_cache = InMemoryCache(max_size_per_item=100) | |
| result = in_memory_cache.check_value_size("a" * 100000000) | |
| assert result is False | |
| def test_in_memory_cache_ttl(): | |
| """ | |
| Check that | |
| - if ttl is not set, it will be set to default ttl | |
| - if object expires, the ttl is also removed | |
| """ | |
| in_memory_cache = InMemoryCache() | |
| in_memory_cache.set_cache(key="my-fake-key", value="my-fake-value", ttl=10) | |
| initial_ttl_time = in_memory_cache.ttl_dict["my-fake-key"] | |
| assert initial_ttl_time is not None | |
| in_memory_cache.set_cache(key="my-fake-key", value="my-fake-value-2", ttl=10) | |
| new_ttl_time = in_memory_cache.ttl_dict["my-fake-key"] | |
| assert new_ttl_time == initial_ttl_time # ttl should not be updated | |
| ## On object expiration, the ttl should be removed | |
| in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value", ttl=1) | |
| new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"] | |
| assert new_ttl_time is not None | |
| time.sleep(1) | |
| cached_obj = in_memory_cache.get_cache(key="new-fake-key") | |
| new_ttl_time = in_memory_cache.ttl_dict.get("new-fake-key") | |
| assert new_ttl_time is None | |
| def test_in_memory_cache_ttl_allow_override(): | |
| """ | |
| Check that | |
| - if ttl is not set, it will be set to default ttl | |
| - if object expires, the ttl is also removed | |
| """ | |
| in_memory_cache = InMemoryCache() | |
| ## On object expiration, but no get_cache, the override should be allowed | |
| in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value", ttl=1) | |
| initial_ttl_time = in_memory_cache.ttl_dict["new-fake-key"] | |
| assert initial_ttl_time is not None | |
| time.sleep(1) | |
| in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value-2", ttl=1) | |
| new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"] | |
| assert new_ttl_time is not None | |
| assert new_ttl_time != initial_ttl_time | |