Add tests for keyed buckets

Author: k-yomo

src/aggregation/agg_req.rs

@@ -271,7 +271,7 @@ mod tests {
                         (7f64..20f64).into(),
                         (20f64..f64::MAX).into(),
                     ],
-                    ..Default::default()
+                    keyed: Some(true),
                 }),
                 sub_aggregation: Default::default(),
             }),
@@ -298,7 +298,8 @@ mod tests {
         {
           "from": 20.0
         }
-      ]
+      ],
+      "keyed": true
     }
   }
 }"#;

src/aggregation/bucket/histogram/histogram.rs

@@ -1398,4 +1398,46 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn histogram_keyed_buckets_test() -> crate::Result<()> {
+        let index = get_test_index_with_num_docs(false, 100)?;
+
+        let agg_req: Aggregations = vec![(
+            "histogram".to_string(),
+            Aggregation::Bucket(BucketAggregation {
+                bucket_agg: BucketAggregationType::Histogram(HistogramAggregation {
+                    field: "score_f64".to_string(),
+                    interval: 50.0,
+                    keyed: Some(true),
+                    ..Default::default()
+                }),
+                sub_aggregation: Default::default(),
+            }),
+        )]
+            .into_iter()
+            .collect();
+
+        let res = exec_request(agg_req, &index)?;
+
+        assert_eq!(
+            res,
+            json!({
+                "histogram": {
+                    "buckets": {
+                        "0": {
+                            "key": 0.0,
+                            "doc_count": 50
+                        },
+                        "50": {
+                            "key": 50.0,
+                            "doc_count": 50
+                        }
+                    }
+                }
+            })
+        );
+
+        Ok(())
+    }
 }

src/aggregation/bucket/range.rs

@@ -457,6 +457,49 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn range_keyed_buckets_test() -> crate::Result<()> {
+        let index = get_test_index_with_num_docs(false, 100)?;
+
+        let agg_req: Aggregations = vec![(
+            "range".to_string(),
+            Aggregation::Bucket(BucketAggregation {
+                bucket_agg: BucketAggregationType::Range(RangeAggregation {
+                    field: "fraction_f64".to_string(),
+                    ranges: vec![(0f64..0.1f64).into(), (0.1f64..0.2f64).into()],
+                    keyed: Some(true),
+                }),
+                sub_aggregation: Default::default(),
+            }),
+        )]
+            .into_iter()
+            .collect();
+
+        let collector = AggregationCollector::from_aggs(agg_req, None);
+
+        let reader = index.reader()?;
+        let searcher = reader.searcher();
+        let agg_res = searcher.search(&AllQuery, &collector).unwrap();
+
+        let res: Value = serde_json::from_str(&serde_json::to_string(&agg_res)?)?;
+
+        assert_eq!(
+            res,
+            json!({
+                "range": {
+                    "buckets": {
+                        "*-0": { "key": "*-0", "doc_count": 0, "to": 0.0},
+                        "0-0.1": {"key": "0-0.1", "doc_count": 10, "from": 0.0, "to": 0.1},
+                        "0.1-0.2": {"key": "0.1-0.2", "doc_count": 10, "from": 0.1, "to": 0.2},
+                        "0.2-*": {"key": "0.2-*", "doc_count": 80, "from": 0.2},
+                    }
+                }
+            })
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn bucket_test_extend_range_hole() {
         let buckets = vec![(10f64..20f64).into(), (30f64..40f64).into()];