chore: test retain_at_most_two_hits_per_file (#3418)

crates/tabby-common/src/api/code.rs

@@ -7,21 +7,21 @@ pub struct CodeSearchResponse {
     pub hits: Vec<CodeSearchHit>,
 }
 
-#[derive(Default, Clone)]
+#[derive(Default, Clone, PartialEq, Debug)]
 pub struct CodeSearchHit {
     pub scores: CodeSearchScores,
     pub doc: CodeSearchDocument,
 }
 
-#[derive(Default, Clone)]
+#[derive(Default, Clone, PartialEq, Debug)]
 pub struct CodeSearchScores {
     /// Reciprocal rank fusion score: https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html
     pub rrf: f32,
     pub bm25: f32,
     pub embedding: f32,
 }
 
-#[derive(Builder, Default, Clone)]
+#[derive(Builder, Default, Clone, PartialEq, Debug)]
 pub struct CodeSearchDocument {
     /// Unique identifier for the file in the repository, stringified SourceFileKey.
     pub file_id: String,

crates/tabby/src/services/code.rs

@@ -263,3 +263,94 @@ impl CodeSearch for CodeSearchService {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_retain_at_most_two_hits_per_file() {
+        let new_hit = |file: &str, body: &str| CodeSearchHit {
+            scores: Default::default(),
+            doc: CodeSearchDocument {
+                file_id: file.to_string(),
+                chunk_id: "chunk1".to_owned(),
+                body: body.to_string(),
+                filepath: "".to_owned(),
+                git_url: "".to_owned(),
+                language: "".to_owned(),
+                start_line: 0,
+            },
+        };
+
+        let cases = vec![
+            (vec![], vec![]),
+            (
+                vec![new_hit("file1", "body1")],
+                vec![new_hit("file1", "body1")],
+            ),
+            (
+                vec![new_hit("file1", "body1"), new_hit("file1", "body2")],
+                vec![new_hit("file1", "body1"), new_hit("file1", "body2")],
+            ),
+            (
+                vec![
+                    new_hit("file1", "body1"),
+                    new_hit("file1", "body2"),
+                    new_hit("file1", "body3"),
+                ],
+                vec![new_hit("file1", "body1"), new_hit("file1", "body2")],
+            ),
+            (
+                vec![
+                    new_hit("file1", "body1"),
+                    new_hit("file1", "body2"),
+                    new_hit("file1", "body3"),
+                    new_hit("file2", "body4"),
+                ],
+                vec![
+                    new_hit("file1", "body1"),
+                    new_hit("file1", "body2"),
+                    new_hit("file2", "body4"),
+                ],
+            ),
+            (
+                vec![
+                    new_hit("file1", "body1"),
+                    new_hit("file1", "body2"),
+                    new_hit("file1", "body3"),
+                    new_hit("file2", "body4"),
+                    new_hit("file2", "body5"),
+                ],
+                vec![
+                    new_hit("file1", "body1"),
+                    new_hit("file1", "body2"),
+                    new_hit("file2", "body4"),
+                    new_hit("file2", "body5"),
+                ],
+            ),
+            (
+                vec![
+                    new_hit("file1", "body1"),
+                    new_hit("file1", "body2"),
+                    new_hit("file1", "body3"),
+                    new_hit("file2", "body4"),
+                    new_hit("file2", "body5"),
+                    new_hit("file2", "body6"),
+                ],
+                vec![
+                    new_hit("file1", "body1"),
+                    new_hit("file1", "body2"),
+                    new_hit("file2", "body4"),
+                    new_hit("file2", "body5"),
+                ],
+            ),
+        ];
+
+        for (input, expected) in cases {
+            let mut input = input;
+            retain_at_most_two_hits_per_file(&mut input);
+            assert_eq!(input, expected);
+        }
+    }
+}