mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-24 11:40:20 +00:00
Krish Dholakia
ef42461c1e
* test: add __init__.py files * refactor: rename test folder to avoid naming conflict * test: update workflows * test: update tests * test: update imports * test: update tests * test: remove unused import * ci(test-litellm.yml): add pytest retry to github workflow * test: fix test
91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
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
|