test: fix flaky tests with SubAgent client builder

Author: rubenruizdegauna

agent-control/src/opamp/client_builder.rs

@@ -100,6 +100,8 @@ pub(crate) mod tests {
         Client, ClientResult, NotStartedClient, NotStartedClientResult, StartedClient,
         StartedClientResult,
     };
+    use std::thread;
+    use std::thread::sleep;
 
     use super::*;
 
@@ -220,5 +222,34 @@ pub(crate) mod tests {
                 .once()
                 .return_once(move |_, _, _| Ok(client));
         }
+
+        // This is a Mock OpAMP Client Builder, which builds the Callbacks and the OpAMP Client
+        // and starts the OpAMP Client thread. This thread owns the callbacks, and these publish
+        // into the OpAMP Publisher <-- Sub Agent OpAMP Consumer
+        // Sub Agent OpAMP Consumer consumes the OpAMP events.
+        // Using the Mock, makes the OpAMP publisher to be dropped, as it's part of the expectations
+        // Until we refactor these tests (and find a better solution for this pattern) we'll
+        // spawn a thread with the publisher, just not to be dropped in the test
+        // TL;DR: Let the OpAMP Publisher leave for Duration
+        pub fn should_build_and_start_and_run(
+            &mut self,
+            agent_id: AgentID,
+            start_settings: StartSettings,
+            client: MockStartedOpAMPClientMock<C>,
+            run_for: Duration,
+        ) {
+            self.expect_build_and_start()
+                .withf(move |publisher, _sub_agent_id, _start_settings| {
+                    let publisher = publisher.clone();
+                    thread::spawn(move || {
+                        sleep(run_for);
+                        drop(publisher)
+                    });
+                    //
+                    agent_id == _sub_agent_id.clone() && start_settings == *_start_settings
+                })
+                .once()
+                .return_once(move |_, _, _| Ok(client));
+        }
     }
 }

agent-control/src/sub_agent/on_host/builder.rs

@@ -240,6 +240,9 @@ mod tests {
         AgentDescription, DescriptionValueType, StartSettings,
     };
     use std::collections::HashMap;
+    use std::thread;
+    use std::thread::sleep;
+    use std::time::Duration;
     use tracing_test::traced_test;
 
     // TODO: tests below are testing not only the builder but also the sub-agent start/stop behavior.
@@ -277,10 +280,13 @@ mod tests {
         started_client.should_stop(1);
 
         // Infra Agent OpAMP no final stop nor health, just after stopping on reload
-        opamp_builder.should_build_and_start(
+        // TODO: We should discussed if this is a valid approach once we refactor the unit tests
+        // Build an OpAMP Client and let it run so the publisher is not dropped
+        opamp_builder.should_build_and_start_and_run(
             sub_agent_id.clone(),
             start_settings_infra,
             started_client,
+            Duration::from_millis(10),
         );
 
         let mut instance_id_getter = MockInstanceIDGetterMock::new();
@@ -320,12 +326,15 @@ mod tests {
             Arc::new(signature_validator),
         );
 
-        on_host_builder
+        let agent = on_host_builder
             .build(sub_agent_id, &sub_agent_config, opamp_publisher)
             .unwrap()
-            .run()
-            .stop()
-            .unwrap();
+            .run();
+
+        // Force OpAMP Publisher drop error
+        // sleep(Duration::from_millis(30));
+
+        agent.stop().unwrap();
     }
 
     #[traced_test]
@@ -375,10 +384,13 @@ mod tests {
         started_client.should_update_effective_config(1);
         started_client.should_stop(1);
 
-        opamp_builder.should_build_and_start(
+        // TODO: We should discussed if this is a valid approach once we refactor the unit tests
+        // Build an OpAMP Client and let it run so the publisher is not dropped
+        opamp_builder.should_build_and_start_and_run(
             sub_agent_id.clone(),
             start_settings_infra,
             started_client,
+            Duration::from_millis(10),
         );
 
         effective_agent_assembler.should_assemble_agent(
@@ -412,6 +424,10 @@ mod tests {
             .build(sub_agent_id, &sub_agent_config, opamp_publisher)
             .expect("Subagent build should be OK");
         let started_sub_agent = sub_agent.run(); // Running the sub-agent should report the failed configuration.
+
+        // Force OpAMP Publisher drop error
+        // sleep(Duration::from_millis(20));
+
         started_sub_agent.stop().unwrap();
     }