add top_k=2 at test_code

Marker-Inc-Korea · Apr 22, 2024 · 710e2af · 710e2af
1 parent 15d9806
commit 710e2af
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/autorag/nodes/passageaugmenter/test_prev_next_augmenter.py b/tests/autorag/nodes/passageaugmenter/test_prev_next_augmenter.py
@@ -33,12 +33,12 @@ def test_prev_next_augmenter_multi_passages():
 
 
 def test_prev_next_augmenter_node():
-    result_df = prev_next_augmenter(project_dir=project_dir, previous_result=previous_result, mode='next', top_k=3)
+    result_df = prev_next_augmenter(project_dir=project_dir, previous_result=previous_result, mode='next', top_k=2)
     contents = result_df["retrieved_contents"].tolist()
     ids = result_df["retrieved_ids"].tolist()
     scores = result_df["retrieve_scores"].tolist()
     assert len(contents) == len(ids) == len(scores) == 2
-    assert len(contents[0]) == len(ids[0]) == len(scores[0]) == 3
+    assert len(contents[0]) == len(ids[0]) == len(scores[0]) == 2
     for content_list, id_list, score_list in zip(contents, ids, scores):
         for i, (content, _id, score) in enumerate(zip(content_list, id_list, score_list)):
             assert isinstance(content, str)