[chore]: Annotation Added for Tests LLamaCpp

langchain-ai · Dec 20, 2024 · 1572f3b · 1572f3b
1 parent d301422
commit 1572f3b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/libs/community/tests/unit_tests/embeddings/test_llamacpp.py b/libs/community/tests/unit_tests/embeddings/test_llamacpp.py
@@ -15,12 +15,12 @@ def mock_llama_client():
         yield mock_client
 
 
-def test_initialization(mock_llama_client):
+def test_initialization(mock_llama_client: MagicMock) -> None:
     embeddings = LlamaCppEmbeddings(client=mock_llama_client)
     assert embeddings.client is not None
 
 
-def test_embed_documents(mock_llama_client):
+def test_embed_documents(mock_llama_client: MagicMock) -> None:
     mock_llama_client.create_embedding.return_value = {
         "data": [{"embedding": [[0.1, 0.2, 0.3]]}, {"embedding": [[0.4, 0.5, 0.6]]}]
     }
@@ -31,7 +31,7 @@ def test_embed_documents(mock_llama_client):
     assert result == expected
 
 
-def test_embed_query(mock_llama_client):
+def test_embed_query(mock_llama_client: MagicMock) -> None:
     mock_llama_client.embed.return_value = [[0.1, 0.2, 0.3]]
     embeddings = LlamaCppEmbeddings(client=mock_llama_client)
     result = embeddings.embed_query("Sample query")