diff --git a/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py b/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py index 8a1be64365a..50869b985f2 100644 --- a/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py +++ b/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py @@ -1,9 +1,11 @@ """Test Neo4jVector functionality.""" import os +from math import isclose from typing import Any, Dict, List, cast from langchain_core.documents import Document +from yaml import safe_load from langchain_community.graphs import Neo4jGraph from langchain_community.vectorstores.neo4j_vector import ( @@ -217,12 +219,21 @@ def test_neo4jvector_relevance_score() -> None: ) output = docsearch.similarity_search_with_relevance_scores("foo", k=3) - assert output == [ + expected_output = [ (Document(page_content="foo", metadata={"page": "0"}), 1.0), (Document(page_content="bar", metadata={"page": "1"}), 0.9998376369476318), (Document(page_content="baz", metadata={"page": "2"}), 0.9993523359298706), ] + # Check if the length of the outputs matches + assert len(output) == len(expected_output) + + # Check if each document and its relevance score is close to the expected value + for (doc, score), (expected_doc, expected_score) in zip(output, expected_output): + assert doc.page_content == expected_doc.page_content + assert doc.metadata == expected_doc.metadata + assert isclose(score, expected_score, rel_tol=1e-5) + drop_vector_indexes(docsearch) @@ -779,8 +790,16 @@ def test_retrieval_dictionary() -> None: ) ) ] + output = docsearch.similarity_search("Foo", k=1) - assert output == expected_output + + def parse_document(doc: Document) -> Any: + return safe_load(doc.page_content) + + parsed_expected = [parse_document(doc) for doc in expected_output] + parsed_output = [parse_document(doc) for doc in output] + + assert parsed_output == parsed_expected drop_vector_indexes(docsearch)