From 92660451a0e4e44bece6ccdd70404c4b4ad24ee3 Mon Sep 17 00:00:00 2001
From: Tomaz Bratanic <bratanic.tomaz@gmail.com>
Date: Sat, 7 Dec 2024 23:10:49 +0100
Subject: [PATCH] fix test

---
 libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py b/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py
index 82b0153..837cb79 100644
--- a/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py
+++ b/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py
@@ -203,8 +203,9 @@ def test_converting_to_yaml() -> None:
 def test_get_search_index_query_hybrid_node_neo4j_5_23_above() -> None:
     expected_query = (
         "CALL () { "
-        "CALL db.index.vector.queryNodes($index, $k, $embedding) "
+        "CALL db.index.vector.queryNodes($index, $k * $ef, $embedding) "
         "YIELD node, score "
+        "WITH node, score LIMIT $k "
         "WITH collect({node:node, score:score}) AS nodes, max(score) AS max "
         "UNWIND nodes AS n "
         "RETURN n.node AS node, (n.score / max) AS score UNION "
@@ -225,8 +226,9 @@ def test_get_search_index_query_hybrid_node_neo4j_5_23_above() -> None:
 def test_get_search_index_query_hybrid_node_neo4j_5_23_below() -> None:
     expected_query = (
         "CALL { "
-        "CALL db.index.vector.queryNodes($index, $k, $embedding) "
+        "CALL db.index.vector.queryNodes($index, $k * $ef, $embedding) "
         "YIELD node, score "
+        "WITH node, score LIMIT $k "
         "WITH collect({node:node, score:score}) AS nodes, max(score) AS max "
         "UNWIND nodes AS n "
         "RETURN n.node AS node, (n.score / max) AS score UNION "