WIP: encapsulate snapshot assertion and URI gen

trace-utils/tests/test_send_data.rs

@@ -24,25 +24,25 @@ mod tests {
     use ddcommon::Endpoint;
 
     #[derive(Debug)]
-    pub struct DatadogTestAgent {
+    pub struct DatadogTestAgentContainer {
         mounts: Vec<Mount>,
     }
 
-    impl Default for DatadogTestAgent {
+    impl Default for DatadogTestAgentContainer {
         fn default() -> Self {
             let mount = Mount::bind_mount(
-                DatadogTestAgent::calculate_snapshot_absolute_path(),
+                DatadogTestAgentContainer::calculate_snapshot_absolute_path(),
                 "/snapshots",
             )
             .with_access_mode(AccessMode::ReadWrite);
 
-            DatadogTestAgent {
+            DatadogTestAgentContainer {
                 mounts: vec![mount],
             }
         }
     }
 
-    impl Image for DatadogTestAgent {
+    impl Image for DatadogTestAgentContainer {
         type Args = Vec<String>;
 
         fn name(&self) -> String {
@@ -55,8 +55,15 @@ mod tests {
 
         fn ready_conditions(&self) -> Vec<WaitFor> {
             vec![
-                WaitFor::message_on_stderr("INFO:ddapm_test_agent.agent:Trace request stall seconds setting set to 0.0."),
-                WaitFor::Duration{length: Duration::from_secs(5)}
+                WaitFor::message_on_stderr(
+                    "INFO:ddapm_test_agent.agent:Trace request stall seconds setting set to 0.0.",
+                ),
+                // Adding this wait because github runners are inconsistent and sometimes they
+                // reset the network TODO: EK - Is this necessary? Can we improve
+                // "ready" signal on test-agent?
+                WaitFor::Duration {
+                    length: Duration::from_secs(3),
+                },
             ]
         }
 
@@ -68,7 +75,7 @@ mod tests {
         }
     }
 
-    impl DatadogTestAgent {
+    impl DatadogTestAgentContainer {
         fn calculate_snapshot_absolute_path() -> String {
             // Get the path of the current executable
             let exe_path = env::current_exe().expect("Failed to get current executable");
@@ -100,21 +107,70 @@ mod tests {
         }
     }
 
+    pub struct DatadogTestAgent {
+        container: ContainerAsync<DatadogTestAgentContainer>,
+    }
+
+    impl DatadogTestAgent {
+        // TODO: EK - can't have async default. Is there a better solution?
+        pub async fn new() -> Self {
+            DatadogTestAgent {
+                container: DatadogTestAgentContainer::default()
+                    .start()
+                    .await
+                    .expect("Unable to start DatadogTestAgent"),
+            }
+        }
+
+        async fn get_base_uri_string(&self) -> String {
+            let container_host = self.container.get_host().await.unwrap().to_string();
+            let container_port = self.container.get_host_port_ipv4(8126).await.unwrap();
+
+            format!("http://{}:{}", container_host, container_port)
+        }
+
+        pub async fn get_uri_for_endpoint(
+            &self,
+            endpoint: &str,
+            snapshot_token: Option<&str>,
+        ) -> Uri {
+            let base_uri_string = self.get_base_uri_string().await;
+            let uri_string = match snapshot_token {
+                Some(token) => format!(
+                    "{}/{}?test_session_token={}",
+                    base_uri_string, endpoint, token
+                ),
+                None => format!("{}/{}", base_uri_string, endpoint),
+            };
+
+            Uri::from_str(&uri_string).expect("Invalid URI")
+        }
+
+        pub async fn assert_snapshot(&self, snapshot_token: &str) {
+            let client = Client::new();
+            let uri = self
+                .get_uri_for_endpoint("test/session/snapshot", Some(snapshot_token))
+                .await;
+            let res = client.get(uri).await.expect("Request failed");
+            let status_code = res.status();
+            let body_bytes = to_bytes(res.into_body()).await.expect("Read failed");
+            let body_string = String::from_utf8(body_bytes.to_vec()).expect("Conversion failed");
+
+            assert_eq!(
+                status_code, 200,
+                "Expected status 200, but got {}. Response body: {}",
+                status_code, body_string
+            );
+        }
+    }
+
     #[cfg_attr(miri, ignore)]
     #[tokio::test]
+    // TODO: EK - This test is just for debugging purposes. Remove before merging.
     async fn connect_to_test_agent_test() {
-        let container = DatadogTestAgent::default()
-            .start()
-            .await
-            .expect("Unable to start DatadogTestAgent");
-        let container_host = container.get_host().await.unwrap().to_string();
-        let container_port = container.get_host_port_ipv4(8126).await.unwrap();
-
         let client = Client::new();
-
-        let uri_string = format!("http://{}:{}/info", container_host, container_port);
-        let uri = Uri::from_str(&uri_string).expect("Invalid URI");
-
+        let test_agent = DatadogTestAgent::new().await;
+        let uri = test_agent.get_uri_for_endpoint("info", None).await;
         let res = client.get(uri).await.expect("Request failed");
         let body_bytes = to_bytes(res.into_body()).await.expect("Read failed");
         let body_string = String::from_utf8(body_bytes.to_vec()).expect("Conversion failed");
@@ -128,12 +184,7 @@ mod tests {
     #[cfg_attr(miri, ignore)]
     #[tokio::test]
     async fn compare_v04_trace_snapshot_test() {
-        let container = DatadogTestAgent::default()
-            .start()
-            .await
-            .expect("Unable to start DatadogTestAgent");
-        let container_host = container.get_host().await.unwrap().to_string();
-        let container_port = container.get_host_port_ipv4(8126).await.unwrap();
+        let test_agent = DatadogTestAgent::new().await;
 
         let header_tags = TracerHeaderTags {
             lang: "test-lang",
@@ -146,15 +197,10 @@ mod tests {
             client_computed_stats: false,
         };
 
-        let agent_trace_endpoint_string_uri = format!(
-            "http://{}:{}/v0.4/traces?test_session_token=compare_v04_trace_snapshot_test",
-            container_host, container_port
-        );
-        let agent_trace_endpoint_uri =
-            Uri::from_str(&agent_trace_endpoint_string_uri).expect("Invalid URI");
-
         let endpoint = Endpoint {
-            url: agent_trace_endpoint_uri,
+            url: test_agent
+                .get_uri_for_endpoint("v0.4/traces", Some("compare_v04_trace_snapshot_test"))
+                .await,
             api_key: None,
         };
 
@@ -169,21 +215,8 @@ mod tests {
 
         let _result = data.send().await;
 
-        let client = Client::new();
-        let uri_string = format!(
-            "http://{}:{}/test/session/snapshot?test_session_token=compare_v04_trace_snapshot_test",
-            container_host, container_port
-        );
-        let uri = Uri::from_str(&uri_string).expect("Invalid URI");
-        let res = client.get(uri).await.expect("Request failed");
-        let status_code = res.status();
-        let body_bytes = to_bytes(res.into_body()).await.expect("Read failed");
-        let body_string = String::from_utf8(body_bytes.to_vec()).expect("Conversion failed");
-
-        assert_eq!(
-            status_code, 200,
-            "Expected status 200, but got {}. Response body: {}",
-            status_code, body_string
-        );
+        test_agent
+            .assert_snapshot("compare_v04_trace_snapshot_test")
+            .await;
     }
 }