test find_top_k also

1 files changed, 21 insertions, 1 deletions

diff --git a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
index 081361b2fdc..1965d58c0e7 100644
--- a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
+++ b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
@@ -59,7 +59,7 @@ public:
     {
         vectors.set(1, {2, 2}).set(2, {3, 2}).set(3, {2, 3})
                .set(4, {1, 2}).set(5, {8, 3}).set(6, {7, 2})
-               .set(7, {3, 5});
+               .set(7, {3, 5}).set(8, {0, 3}).set(9, {4, 5});
     }
     void init(bool heuristic_select_neighbors) {
         index = std::make_unique<HnswIndex>(vectors, distance_func, level_generator,
@@ -86,6 +86,16 @@ public:
         ASSERT_EQ(exp_levels.size(), act_node.size());
         EXPECT_EQ(exp_levels, act_node.levels());
     }
+    void expect_top_3(uint32_t docid, std::vector<uint32_t> exp_hits) {
+        auto rv = index->find_top_k(vectors.get_vector(docid), 3);
+        size_t idx = 0;
+        for (const auto & hit : rv) {
+            // fprintf(stderr, "found docid %u dist %.1f\n", hit.docid, hit.distance);
+            if (idx < exp_hits.size()) {
+                EXPECT_EQ(hit.docid, exp_hits[idx++]);
+            }
+        }
+    }
 };
 
 
@@ -134,6 +144,16 @@ TEST_F(HnswIndexTest, 2d_vectors_inserted_in_level_0_graph_with_simple_select_ne
     expect_level_0(5, {2, 3, 6});
     expect_level_0(6, {2, 5});
     expect_level_0(7, {2, 3});
+
+    expect_top_3(1, {1});
+    expect_top_3(2, {2, 1, 3});
+    expect_top_3(3, {3});
+    expect_top_3(4, {4, 1, 3});
+    expect_top_3(5, {5, 6, 2});
+    expect_top_3(6, {6, 5, 2});
+    expect_top_3(7, {7, 3, 2});
+    expect_top_3(8, {4, 3, 1});
+    expect_top_3(9, {7, 3, 2});
 }
 
 TEST_F(HnswIndexTest, 2d_vectors_inserted_and_removed)