From 654cd33fc7b2e11a89346b42102d8790c7cdc98f Mon Sep 17 00:00:00 2001
From: Ben Frederickson <ben@benfrederickson.com>
Date: Wed, 17 Jul 2024 11:08:20 -0700
Subject: [PATCH] Mark the rust brute force unittest as flaky

The rust `brute_force::tests::test_l2` test can sometimes randomly
fail (https://github.com/rapidsai/cuvs/issues/127).

While we should figure out whats causing the test failure, but in
the meantime - this PR marks the test as flaky using the `mark_flaky_tests`
crate, which will retry the test 3 times and pass it if any
of those runs work.
---
 rust/cuvs/Cargo.toml          | 1 +
 rust/cuvs/src/brute_force.rs  | 3 ++-
 rust/cuvs/src/distance/mod.rs | 3 ---
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/rust/cuvs/Cargo.toml b/rust/cuvs/Cargo.toml
index d69e86f9b..8ae4374c1 100644
--- a/rust/cuvs/Cargo.toml
+++ b/rust/cuvs/Cargo.toml
@@ -14,3 +14,4 @@ ndarray = "0.15"
 
 [dev-dependencies]
 ndarray-rand = "0.14"
+mark-flaky-tests = "1"
diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs
index 975a4ddc3..c153359b8 100644
--- a/rust/cuvs/src/brute_force.rs
+++ b/rust/cuvs/src/brute_force.rs
@@ -106,6 +106,7 @@ mod tests {
     use ndarray::s;
     use ndarray_rand::rand_distr::Uniform;
     use ndarray_rand::RandomExt;
+    use mark_flaky_tests::flaky;
 
     fn test_bfknn(metric: DistanceType) {
         let res = Resources::new().unwrap();
@@ -174,7 +175,7 @@ mod tests {
     }
 */
 
-    #[test]
+    #[flaky]
     fn test_l2() {
         test_bfknn(DistanceType::L2Expanded);
     }
diff --git a/rust/cuvs/src/distance/mod.rs b/rust/cuvs/src/distance/mod.rs
index 0377f3e32..566bf0129 100644
--- a/rust/cuvs/src/distance/mod.rs
+++ b/rust/cuvs/src/distance/mod.rs
@@ -15,8 +15,6 @@
  */
 
 
-use std::io::{stderr, Write};
-
 use crate::distance_type::DistanceType;
 use crate::dlpack::ManagedTensor;
 use crate::error::{check_cuvs, Result};
@@ -55,7 +53,6 @@ pub fn pairwise_distance(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use ndarray::s;
     use ndarray_rand::rand_distr::Uniform;
     use ndarray_rand::RandomExt;