diff --git a/libs/community/tests/unit_tests/embeddings/test_llamacpp.py b/libs/community/tests/unit_tests/embeddings/test_llamacpp.py index 423079d1336ad..f7b0bfc62f264 100644 --- a/libs/community/tests/unit_tests/embeddings/test_llamacpp.py +++ b/libs/community/tests/unit_tests/embeddings/test_llamacpp.py @@ -6,7 +6,7 @@ @pytest.fixture -def mock_llama_client(): +def mock_llama_client() -> MagicMock: with patch( "langchain_community.embeddings.llamacpp.LlamaCppEmbeddings" ) as MockLlama: @@ -16,7 +16,7 @@ def mock_llama_client(): def test_initialization(mock_llama_client: MagicMock) -> None: - embeddings = LlamaCppEmbeddings(client=mock_llama_client) + embeddings = LlamaCppEmbeddings(client=mock_llama_client) # type: ignore[call-arg] assert embeddings.client is not None @@ -24,7 +24,7 @@ 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]]}] } - embeddings = LlamaCppEmbeddings(client=mock_llama_client) + embeddings = LlamaCppEmbeddings(client=mock_llama_client) # type: ignore[call-arg] texts = ["Hello world", "Test document"] result = embeddings.embed_documents(texts) expected = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] @@ -33,7 +33,7 @@ def test_embed_documents(mock_llama_client: MagicMock) -> None: 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) + embeddings = LlamaCppEmbeddings(client=mock_llama_client) # type: ignore[call-arg] result = embeddings.embed_query("Sample query") expected = [0.1, 0.2, 0.3] assert result == expected