Fix flakiness on cancellation token (#977)

* fix(data-pipeline): fix flakiness on cancellation token

* fix(stats-exporter): remove duplicate check

Author: VianneyRuhlmann

data-pipeline/src/stats_exporter.rs

@@ -306,7 +306,7 @@ mod tests {
     async fn test_run() {
         let server = MockServer::start_async().await;
 
-        let mock = server
+        let mut mock = server
             .mock_async(|when, then| {
                 when.method(POST)
                     .header("Content-type", "application/msgpack")
@@ -332,17 +332,18 @@ mod tests {
         tokio::time::sleep(BUCKETS_DURATION + Duration::from_secs(1)).await;
         // Resume time to sleep while the stats are being sent
         tokio::time::resume();
-        tokio::time::sleep(Duration::from_millis(100)).await;
-
-        mock.assert_async().await;
+        assert!(
+            poll_for_mock_hit(&mut mock, 10, 100, 1, false).await,
+            "Expected max retry attempts"
+        );
     }
 
     #[cfg_attr(miri, ignore)]
     #[tokio::test]
     async fn test_cancellation_token() {
         let server = MockServer::start_async().await;
 
-        let mock = server
+        let mut mock = server
             .mock_async(|when, then| {
                 when.method(POST)
                     .header("Content-type", "application/msgpack")
@@ -368,8 +369,10 @@ mod tests {
         });
         // Cancel token to trigger force flush
         cancellation_token.cancel();
-        tokio::time::sleep(Duration::from_millis(500)).await;
 
-        mock.assert_async().await;
+        assert!(
+            poll_for_mock_hit(&mut mock, 10, 100, 1, false).await,
+            "Expected max retry attempts"
+        );
     }
 }