Move poll_for_mock_hit test helper function to test-trace-utils

sidecar/Cargo.toml

@@ -76,6 +76,7 @@ version = "0.51.0"
 libc = { version = "0.2" }
 tempfile = { version = "3.3" }
 httpmock = "0.7.0"
+datadog-trace-test-utils = { path = "../trace-test-utils" }
 
 [target.'cfg(not(target_arch = "x86"))'.dependencies]
 simd-json = "0.13.8"

sidecar/src/service/tracing/trace_flusher.rs

@@ -301,36 +301,11 @@ impl TraceFlusher {
 mod tests {
     use super::*;
     use datadog_trace_protobuf::pb;
+    use datadog_trace_test_utils::poll_for_mock_hit;
     use datadog_trace_utils::trace_utils::TracerHeaderTags;
-    use httpmock::{Mock, MockServer};
+    use httpmock::MockServer;
     use std::sync::Arc;
 
-    // This function will poll the mock server for "hits" until the expected number of hits is
-    // observed. Then it will delete the mock. In its current form it may not correctly report if
-    // more than the asserted number of hits occurred. More attempts at lower sleep intervals is
-    // preferred to reduce flakiness and test runtime.
-    async fn poll_for_mock_hit(
-        mock: &mut Mock<'_>,
-        poll_attempts: i32,
-        sleep_interval_ms: u64,
-        expected_hits: usize,
-    ) -> bool {
-        let mut mock_hit = mock.hits_async().await == expected_hits;
-
-        let mut mock_observations_remaining = poll_attempts;
-
-        while !mock_hit {
-            tokio::time::sleep(Duration::from_millis(sleep_interval_ms)).await;
-            mock_hit = mock.hits_async().await == expected_hits;
-            mock_observations_remaining -= 1;
-            if mock_observations_remaining == 0 || mock_hit {
-                break;
-            }
-        }
-
-        mock_hit
-    }
-
     fn create_send_data(size: usize, target_endpoint: &Endpoint) -> SendData {
         let tracer_header_tags = TracerHeaderTags::default();
 
@@ -386,12 +361,12 @@ mod tests {
         trace_flusher.enqueue(send_data_1);
         trace_flusher.enqueue(send_data_2);
 
-        assert!(poll_for_mock_hit(&mut mock, 10, 150, 0).await);
+        assert!(poll_for_mock_hit(&mut mock, 10, 150, 0, false).await);
 
         // enqueue a trace that exceeds the min force flush size
         trace_flusher.enqueue(send_data_3);
 
-        assert!(poll_for_mock_hit(&mut mock, 25, 100, 1).await);
+        assert!(poll_for_mock_hit(&mut mock, 25, 100, 1, true).await);
     }
 
     #[cfg_attr(miri, ignore)]
@@ -427,7 +402,7 @@ mod tests {
             trace_flusher.interval_ms.load(Ordering::Relaxed) + 1,
         ))
         .await;
-        assert!(poll_for_mock_hit(&mut mock, 25, 100, 1).await);
+        assert!(poll_for_mock_hit(&mut mock, 25, 100, 1, true).await);
     }
 
     #[cfg_attr(miri, ignore)]
@@ -459,6 +434,6 @@ mod tests {
 
         trace_flusher.enqueue(send_data_1);
 
-        assert!(poll_for_mock_hit(&mut mock, 5, 250, 0).await);
+        assert!(poll_for_mock_hit(&mut mock, 5, 250, 0, true).await);
     }
 }

trace-test-utils/Cargo.toml

@@ -11,3 +11,5 @@ license.workspace = true
 [dependencies]
 serde_json = "1.0.117"
 datadog-trace-protobuf = { path = "../trace-protobuf" }
+httpmock = "0.7.0"
+tokio = "1.37.0"

trace-test-utils/src/lib.rs

@@ -2,9 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::collections::HashMap;
+use std::time::Duration;
 
 use datadog_trace_protobuf::pb;
+use httpmock::Mock;
 use serde_json::json;
+use tokio::time::sleep;
 
 pub fn create_test_span(
     trace_id: u64,
@@ -78,3 +81,73 @@ pub fn create_test_json_span(
         }
     )
 }
+/// This is a helper function for observing if a httpmock object has been "hit" the expected number
+/// of times. If not it will perform a tokio::sleep and try again. If `delete_after_hits` is set to
+/// true it will delete the mock. More attempts at lower sleep intervals is preferred to reduce
+/// flakiness and test runtime. This is especially useful when testing async code that may not block
+/// on receiving a response.
+///
+/// # Arguments
+///
+/// * `mock` - A mutable reference to the Mock object.
+/// * `poll_attempts` - The number of times to attempt polling the mock server.
+/// * `sleep_interval_ms` - The interval in milliseconds to sleep between each poll attempt.
+/// * `expected_hits` - The expected number of hits on the mock server.
+/// * `delete_after_hits` - A boolean indicating whether to delete the mock after a hit is observed.
+///
+/// # Returns
+///
+/// * A boolean indicating whether the expected number of hits was observed on the mock.
+///
+/// # Examples
+///
+/// ```
+/// use datadog_trace_test_utils::poll_for_mock_hit;
+/// use httpmock::MockServer;
+///
+/// #[cfg_attr(miri, ignore)]
+/// #[tokio::test]
+/// async fn test_with_poll() {
+///     let server = MockServer::start();
+///
+///     let mut mock = server
+///         .mock_async(|_when, then| {
+///             then.status(202)
+///                 .header("content-type", "application/json")
+///                 .body(r#"{"status":"ok"}"#);
+///         })
+///         .await;
+/// 
+///     // Do something that would trigger a request to the mock server
+/// 
+///     assert!(
+///         poll_for_mock_hit(&mut mock, 10, 100, 1, true).await,
+///         "Expected a request"
+///     );
+/// }
+/// ```
+pub async fn poll_for_mock_hit(
+    mock: &mut Mock<'_>,
+    poll_attempts: i32,
+    sleep_interval_ms: u64,
+    expected_hits: usize,
+    delete_after_hits: bool,
+) -> bool {
+    let mut mock_hit = mock.hits_async().await == expected_hits;
+
+    let mut mock_observations_remaining = poll_attempts;
+
+    while !mock_hit {
+        sleep(Duration::from_millis(sleep_interval_ms)).await;
+        mock_hit = mock.hits_async().await == expected_hits;
+        mock_observations_remaining -= 1;
+        if mock_observations_remaining == 0 || mock_hit {
+            if delete_after_hits {
+                mock.delete();
+            }
+            break;
+        }
+    }
+
+    mock_hit
+}

trace-utils/src/send_data.rs

@@ -389,7 +389,8 @@ impl SendData {
 // TODO: APMSP-1079 - We should have more comprehensive tests for SendData logic beyond retry logic.
 mod tests {
     use super::*;
-    use httpmock::{Mock, MockServer};
+    use datadog_trace_test_utils::poll_for_mock_hit;
+    use httpmock::MockServer;
     use std::time::Duration;
     use tokio::time::Instant;
 
@@ -529,35 +530,6 @@ mod tests {
         );
     }
 
-    // TODO: APMSP-1153 - This function also exists in
-    // sidecar::service::tracing::trace_flusher::tests. It should be moved to a common
-    // trace_test_utils module when it is properly gated to just test dependency.
-    async fn poll_for_mock_hit(
-        mock: &mut Mock<'_>,
-        poll_attempts: i32,
-        sleep_interval_ms: u64,
-        expected_hits: usize,
-        delete_after_hit: bool,
-    ) -> bool {
-        let mut mock_hit = mock.hits_async().await == expected_hits;
-
-        let mut mock_observations_remaining = poll_attempts;
-
-        while !mock_hit {
-            sleep(Duration::from_millis(sleep_interval_ms)).await;
-            mock_hit = mock.hits_async().await == expected_hits;
-            mock_observations_remaining -= 1;
-            if mock_observations_remaining == 0 || mock_hit {
-                if delete_after_hit {
-                    mock.delete();
-                }
-                break;
-            }
-        }
-
-        mock_hit
-    }
-
     // TODO: APMSP-1153 - This function also exists in
     // sidecar::service::tracing::trace_flusher::tests. It should be moved to a common
     // trace_test_utils module when it is properly gated to just test dependency.