From 2c89da258eb916d3a3f5ed86c5e48ec9e1402c63 Mon Sep 17 00:00:00 2001
From: Ezra Singh <singhezra@gmail.com>
Date: Fri, 9 Aug 2024 10:35:05 -0400
Subject: [PATCH] Added test to capture edge cases with unsorted count

---
 src/float/query/nearest_n_within.rs | 65 +++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/src/float/query/nearest_n_within.rs b/src/float/query/nearest_n_within.rs
index 3fc1d06..f8fe7bf 100644
--- a/src/float/query/nearest_n_within.rs
+++ b/src/float/query/nearest_n_within.rs
@@ -230,6 +230,71 @@ mod tests {
         }
     }
 
+    #[test]
+    fn can_query_nearest_n_items_unsorted_max_qty() {
+        let mut tree: KdTree<AX, u32, 4, 4, u32> = KdTree::new();
+
+        let content_to_add: [([AX; 4], u32); 16] = [
+            ([0.9f32, 0.0f32, 0.9f32, 0.0f32], 9),
+            ([0.4f32, 0.5f32, 0.4f32, 0.5f32], 4),
+            ([0.12f32, 0.3f32, 0.12f32, 0.3f32], 12),
+            ([0.7f32, 0.2f32, 0.7f32, 0.2f32], 7),
+            ([0.13f32, 0.4f32, 0.13f32, 0.4f32], 13),
+            ([0.6f32, 0.3f32, 0.6f32, 0.3f32], 6),
+            ([0.2f32, 0.7f32, 0.2f32, 0.7f32], 2),
+            ([0.14f32, 0.5f32, 0.14f32, 0.5f32], 14),
+            ([0.3f32, 0.6f32, 0.3f32, 0.6f32], 3),
+            ([0.10f32, 0.1f32, 0.10f32, 0.1f32], 10),
+            ([0.16f32, 0.7f32, 0.16f32, 0.7f32], 16),
+            ([0.1f32, 0.8f32, 0.1f32, 0.8f32], 1),
+            ([0.15f32, 0.6f32, 0.15f32, 0.6f32], 15),
+            ([0.5f32, 0.4f32, 0.5f32, 0.4f32], 5),
+            ([0.8f32, 0.1f32, 0.8f32, 0.1f32], 8),
+            ([0.11f32, 0.2f32, 0.11f32, 0.2f32], 11),
+        ];
+
+        for (point, item) in content_to_add {
+            tree.add(&point, item);
+        }
+
+        let query_point = [0.78f32, 0.55f32, 0.78f32, 0.55f32];
+
+        let radius = 100.0;
+        let max_qty = 1;
+
+        let result_unsorted: Vec<_> = tree
+            .nearest_n_within::<SquaredEuclidean>(&query_point, radius, max_qty, false)
+            .into_iter()
+            .map(|n| (n.distance, n.item))
+            .collect();
+
+        let result_sorted: Vec<_> = tree
+            .nearest_n_within::<SquaredEuclidean>(&query_point, radius, max_qty, true)
+            .into_iter()
+            .map(|n| (n.distance, n.item))
+            .collect();
+
+        assert_eq!(result_unsorted.len(), max_qty);
+        assert_eq!(result_sorted.len(), max_qty);
+
+        let max_qty = 0;
+
+        let result_unsorted: Vec<_> = tree
+            .nearest_n_within::<SquaredEuclidean>(&query_point, radius, max_qty, false)
+            .into_iter()
+            .map(|n| (n.distance, n.item))
+            .collect();
+
+        let result_sorted: Vec<_> = tree
+            .nearest_n_within::<SquaredEuclidean>(&query_point, radius, max_qty, true)
+            .into_iter()
+            .map(|n| (n.distance, n.item))
+            .collect();
+
+        assert_eq!(result_unsorted.len(), max_qty);
+        assert_eq!(result_sorted.len(), max_qty);
+    }
+
     #[test]
     fn can_query_nearest_n_items_within_radius_unsorted_large_scale() {
         const TREE_SIZE: usize = 100_000;