autogen/test/cache/test_cache.py
Wael Karkoub a9171211c7
Streamline Testing with pytest-cov and pytest Defaults (#2490)
* done

* update docs

* try fix

* update workflows

* undo minor fix

* resolve comments

* adds back pytest-asyncio

* minor fix

* add branch coverage

* restore pip install e.

* test with coverage

* fix mypy

* fix coverage + docker + windows combo

* fix bash command

* formatter

* formatter

* one last fix

* I lied, last fix

* fix

* fix retrieve chat test

* fix windows paths

* change cache seed

* down grade openai version

* fix openai mypy

* better error type

* fix image gen cache test

* fix

* experimenting

* fix lmm

* skip cosmos test

* remove cosmos db

* unused imports

* handle more cosmosdb skips

* fix flaky test
2024-04-29 00:27:34 +00:00

110 lines
3.8 KiB
Python
Executable File

#!/usr/bin/env python3 -m pytest
import unittest
from unittest.mock import ANY, MagicMock, patch
try:
from azure.cosmos import CosmosClient
skip_azure = False
except ImportError:
CosmosClient = object
skip_azure = True
from autogen.cache.cache import Cache
class TestCache(unittest.TestCase):
def setUp(self):
self.redis_config = {
"cache_seed": "test_seed",
"redis_url": "redis://test",
"cache_path_root": ".test_cache",
}
self.cosmos_config = {
"cosmos_db_config": {
"connection_string": "AccountEndpoint=https://example.documents.azure.com:443/;",
"database_id": "autogen_cache",
"container_id": "TestContainer",
"cache_seed": "42",
"client": MagicMock(spec=CosmosClient),
}
}
@patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=MagicMock())
def test_redis_cache_initialization(self, mock_cache_factory):
cache = Cache(self.redis_config)
self.assertIsInstance(cache.cache, MagicMock)
mock_cache_factory.assert_called()
@patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=MagicMock())
@unittest.skipIf(skip_azure, "requires azure.cosmos")
def test_cosmosdb_cache_initialization(self, mock_cache_factory):
cache = Cache(self.cosmos_config)
self.assertIsInstance(cache.cache, MagicMock)
mock_cache_factory.assert_called_with(
seed="42",
redis_url=None,
cache_path_root=None,
cosmosdb_config={
"connection_string": "AccountEndpoint=https://example.documents.azure.com:443/;",
"database_id": "autogen_cache",
"container_id": "TestContainer",
"cache_seed": "42",
"client": ANY,
},
)
def context_manager_common(self, config):
mock_cache_instance = MagicMock()
with patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=mock_cache_instance):
with Cache(config) as cache:
self.assertIsInstance(cache, MagicMock)
mock_cache_instance.__enter__.assert_called()
mock_cache_instance.__exit__.assert_called()
def test_redis_context_manager(self):
self.context_manager_common(self.redis_config)
@unittest.skipIf(skip_azure, "requires azure.cosmos")
def test_cosmos_context_manager(self):
self.context_manager_common(self.cosmos_config)
def get_set_common(self, config):
key = "key"
value = "value"
mock_cache_instance = MagicMock()
with patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=mock_cache_instance):
cache = Cache(config)
cache.set(key, value)
cache.get(key)
mock_cache_instance.set.assert_called_with(key, value)
mock_cache_instance.get.assert_called_with(key, None)
def test_redis_get_set(self):
self.get_set_common(self.redis_config)
@unittest.skipIf(skip_azure, "requires azure.cosmos")
def test_cosmos_get_set(self):
self.get_set_common(self.cosmos_config)
def close_common(self, config):
mock_cache_instance = MagicMock()
with patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=mock_cache_instance):
cache = Cache(config)
cache.close()
mock_cache_instance.close.assert_called()
def test_redis_close(self):
self.close_common(self.redis_config)
@unittest.skipIf(skip_azure, "requires azure.cosmos")
def test_cosmos_close(self):
self.close_common(self.cosmos_config)
if __name__ == "__main__":
unittest.main()