improve coverage of llm

gusye1234 · Sep 25, 2024 · ba04e76 · ba04e76
1 parent 032d041
commit ba04e76
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_openai.py b/tests/test_openai.py
@@ -4,6 +4,20 @@
 from nano_graphrag import _llm
 
 
+def test_get_openai_async_client_instance():
+    with patch("nano_graphrag._llm.AsyncOpenAI") as mock_openai:
+        mock_openai.return_value = "CLIENT"
+        client = _llm.get_openai_async_client_instance()
+    assert client == "CLIENT"
+
+
+def test_get_azure_openai_async_client_instance():
+    with patch("nano_graphrag._llm.AsyncAzureOpenAI") as mock_openai:
+        mock_openai.return_value = "AZURE_CLIENT"
+        client = _llm.get_azure_openai_async_client_instance()
+    assert client == "AZURE_CLIENT"
+
+
 @pytest.fixture
 def mock_openai_client():
     with patch("nano_graphrag._llm.get_openai_async_client_instance") as mock_openai: