fix flakiness

Signed-off-by: Andrei Gherghescu <[email protected]>

Author: andrei-ng

plotly_static/src/lib.rs

@@ -1047,8 +1047,12 @@ impl AsyncStaticExporter {
 
     async fn extract(&mut self, html_content: &str, plot: &PlotData<'_>) -> Result<String> {
         let caps = self.build_webdriver_caps()?;
-        debug!("Use WebDriver and headless browser to export static plot");
+        debug!(
+            "Use WebDriver and headless browser to export static plot (offline_mode={}, port={})",
+            self.offline_mode, self.webdriver_port
+        );
         let webdriver_url = format!("{}:{}", self.webdriver_url, self.webdriver_port);
+        debug!("Connecting to WebDriver at {webdriver_url}");
 
         // Reuse existing client or create new one
         let client = if let Some(ref client) = self.webdriver_client {
@@ -1058,8 +1062,8 @@ impl AsyncStaticExporter {
             debug!("Creating new WebDriver session");
             let new_client = ClientBuilder::native()
                 .capabilities(caps)
-                .connect(&webdriver_url)
-                .await
+                    .connect(&webdriver_url)
+                    .await
                 .with_context(|| "WebDriver session error")?;
             self.webdriver_client = Some(new_client.clone());
             new_client
@@ -1191,20 +1195,30 @@ impl AsyncStaticExporter {
 
 #[cfg(test)]
 mod tests {
-    use std::path::PathBuf;
-    use std::sync::atomic::{AtomicU32, Ordering};
-
     use super::*;
+    use std::path::PathBuf;
 
     fn init() {
         let _ = env_logger::try_init();
     }
 
     // Helper to generate unique ports for parallel tests
-    static PORT_COUNTER: AtomicU32 = AtomicU32::new(4444);
-
     fn get_unique_port() -> u32 {
-        PORT_COUNTER.fetch_add(1, Ordering::SeqCst)
+        use std::sync::atomic::{AtomicU32, Ordering};
+        static PORT_COUNTER: AtomicU32 = AtomicU32::new(4444);
+
+        // Before we used this counter to generate unique ports.
+        // >>> PORT_COUNTER.fetch_add(1, Ordering::SeqCst)
+        // However, sometimes the webdriver process is not stopped immediately
+        // and we get port conflicts.
+        // We try to give some time for other webdriver processes to stop so that we don't
+        // get port conflicts.
+        loop {
+            let p = PORT_COUNTER.fetch_add(1, Ordering::SeqCst);
+            if !webdriver::WebDriver::is_webdriver_running(p) {
+                return p;
+            }
+        }
     }
 
     fn create_test_plot() -> serde_json::Value {

plotly_static/src/webdriver.rs

@@ -32,7 +32,7 @@ const WEBDRIVER_BIN: &str = "chromedriver";
 /// Default WebDriver port
 pub(crate) const WEBDRIVER_PORT: u32 = 4444;
 /// Default WebDriver URL
-pub(crate) const WEBDRIVER_URL: &str = "http://localhost";
+pub(crate) const WEBDRIVER_URL: &str = "http://127.0.0.1";
 
 #[cfg(all(feature = "chromedriver", not(target_os = "windows")))]
 pub(crate) fn chrome_default_caps() -> Vec<&'static str> {