fix flakiness

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

Author: andrei-ng

plotly/src/plot.rs

@@ -1065,12 +1065,10 @@ mod tests {
         assert!(std::fs::remove_file(&dst).is_ok());
     }
 
-    #[cfg(feature = "plotly_static")]
     // Helper to generate unique ports for parallel tests
-    static PORT_COUNTER: AtomicU32 = AtomicU32::new(4444);
-
     #[cfg(feature = "plotly_static")]
     fn get_unique_port() -> u32 {
+        static PORT_COUNTER: AtomicU32 = AtomicU32::new(5544);
         PORT_COUNTER.fetch_add(1, Ordering::SeqCst)
     }
 

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 {
@@ -1079,6 +1083,70 @@ impl AsyncStaticExporter {
         // Open the HTML
         client.goto(&url).await?;
 
+        // Ensure DOM is ready and required elements/scripts are available (Windows CI race)
+        {
+            let start = std::time::Instant::now();
+            let timeout = std::time::Duration::from_secs(10);
+            loop {
+                let state = client
+                    .execute("return document.readyState;", vec![])
+                    .await
+                    .unwrap_or_else(|_| serde_json::Value::Null);
+                if state.as_str().map(|s| s == "complete").unwrap_or(false) {
+                    break;
+                }
+                if start.elapsed() > timeout {
+                    return Err(anyhow!(
+                        "Timeout waiting for document.readyState === 'complete'"
+                    ));
+                }
+                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+            }
+        }
+
+        // Wait for Plotly container element
+        {
+            let start = std::time::Instant::now();
+            let timeout = std::time::Duration::from_secs(10);
+            loop {
+                let has_el = client
+                    .execute(
+                        "return !!document.getElementById('plotly-html-element');",
+                        vec![],
+                    )
+                    .await
+                    .unwrap_or_else(|_| serde_json::Value::Bool(false));
+                if has_el.as_bool().unwrap_or(false) {
+                    break;
+                }
+                if start.elapsed() > timeout {
+                    return Err(anyhow!(
+                        "Timeout waiting for #plotly-html-element to appear in DOM"
+                    ));
+                }
+                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+            }
+        }
+
+        // In online mode, ensure Plotly is loaded
+        if !self.offline_mode {
+            let start = std::time::Instant::now();
+            let timeout = std::time::Duration::from_secs(15);
+            loop {
+                let has_plotly = client
+                    .execute("return !!window.Plotly;", vec![])
+                    .await
+                    .unwrap_or_else(|_| serde_json::Value::Bool(false));
+                if has_plotly.as_bool().unwrap_or(false) {
+                    break;
+                }
+                if start.elapsed() > timeout {
+                    return Err(anyhow!("Timeout waiting for Plotly library to load"));
+                }
+                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+            }
+        }
+
         let (js_script, args) = match plot.format {
             ImageFormat::PDF => {
                 // Always use SVG for PDF export
@@ -1191,20 +1259,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 {
@@ -1370,7 +1448,7 @@ mod tests {
 
         let mut exporter = StaticExporterBuilder::default()
             .spawn_webdriver(true)
-            .webdriver_port(get_unique_port())
+            .webdriver_port(5444)
             .build_async()
             .unwrap();
 

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> {