Merge pull request #596 from ag2ai/fix-optional-imports

Fix failing tests with optional imports
ag2ai · Jan 21, 2025 · 6db520f · 6db520f
2 parents 4605ead + 0a51416
commit 6db520f
Show file tree

Hide file tree

Showing 2 changed files with 218 additions and 227 deletions.
diff --git a/test/cache/test_cache.py b/test/cache/test_cache.py
@@ -15,7 +15,7 @@
 with optional_import_block() as result:
     from azure.cosmos import CosmosClient
 
-skip_azure = not result.is_successful
+skip_azure_cosmos = not result.is_successful
 
 
 class TestCache(unittest.TestCase):
@@ -31,7 +31,7 @@ def setUp(self):
                 "database_id": "autogen_cache",
                 "container_id": "TestContainer",
                 "cache_seed": "42",
-                "client": MagicMock(spec=CosmosClient),
+                "client": MagicMock(spec=CosmosClient) if not skip_azure_cosmos else MagicMock(),
             }
         }
 
@@ -42,7 +42,7 @@ def test_redis_cache_initialization(self, mock_cache_factory):
         mock_cache_factory.assert_called()
 
     @patch("autogen.cache.cache_factory.CacheFactory.cache_factory", return_value=MagicMock())
-    @unittest.skipIf(skip_azure, "requires azure.cosmos")
+    @unittest.skipIf(skip_azure_cosmos, "requires azure.cosmos")
     def test_cosmosdb_cache_initialization(self, mock_cache_factory):
         cache = Cache(self.cosmos_config)
         self.assertIsInstance(cache.cache, MagicMock)
@@ -71,7 +71,7 @@ def context_manager_common(self, config):
     def test_redis_context_manager(self):
         self.context_manager_common(self.redis_config)
 
-    @unittest.skipIf(skip_azure, "requires azure.cosmos")
+    @unittest.skipIf(skip_azure_cosmos, "requires azure.cosmos")
     def test_cosmos_context_manager(self):
         self.context_manager_common(self.cosmos_config)
 
@@ -90,7 +90,7 @@ def get_set_common(self, config):
     def test_redis_get_set(self):
         self.get_set_common(self.redis_config)
 
-    @unittest.skipIf(skip_azure, "requires azure.cosmos")
+    @unittest.skipIf(skip_azure_cosmos, "requires azure.cosmos")
     def test_cosmos_get_set(self):
         self.get_set_common(self.cosmos_config)
 
@@ -104,7 +104,7 @@ def close_common(self, config):
     def test_redis_close(self):
         self.close_common(self.redis_config)
 
-    @unittest.skipIf(skip_azure, "requires azure.cosmos")
+    @unittest.skipIf(skip_azure_cosmos, "requires azure.cosmos")
     def test_cosmos_close(self):
         self.close_common(self.cosmos_config)