reuse session in webdriver

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

Author: andrei-ng

.github/workflows/ci.yml

@@ -107,12 +107,12 @@ jobs:
             -ChromeDriverPath "${{ steps.setup-chrome.outputs.chromedriver-path }}"
       
       # Run tests on non-Windows platforms
-      - name: Run tests (non-Windows)
+      - name: Run tests (${{ matrix.os }}) 
         if: matrix.os != 'windows-latest'
         run: cargo test --verbose --workspace --features plotly_ndarray,plotly_image,static_export_default --exclude plotly_kaleido
 
       # Run tests on Windows with Chrome WebDriver
-      - name: Run tests (Windows)
+      - name: Run tests (${{ matrix.os }}) 
         if: matrix.os == 'windows-latest'
         shell: pwsh
         run: |
@@ -127,12 +127,8 @@ jobs:
           Write-Host "WEBDRIVER_PATH: $env:WEBDRIVER_PATH"
           Write-Host "CHROME_PATH: $env:CHROME_PATH"
           
-          cargo test --verbose --workspace --features plotly_ndarray,plotly_image,chromedriver --exclude plotly_kaleido
+          cargo test --verbose --workspace --features plotly_ndarray,plotly_image,static_export_chromedriver --exclude plotly_kaleido
       
-      # Windows-specific cleanup
-      - if: matrix.os == 'windows-latest'
-        run: gci -recurse -filter "*example*"
-
   code-coverage:
     name: Code Coverage
     runs-on: ubuntu-latest

plotly/src/configuration.rs

@@ -10,7 +10,6 @@ pub enum ImageButtonFormats {
     Webp,
 }
 
-// TODO: should this be behind the plotly-kaleido feature?
 #[serde_with::skip_serializing_none]
 #[derive(Serialize, Debug, Default, Clone)]
 pub struct ToImageButtonOptions {

plotly/src/plot.rs

@@ -831,13 +831,6 @@ mod tests {
     use super::*;
     use crate::Scatter;
 
-    // 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)
-    }
-
     fn create_test_plot() -> Plot {
         let trace1 = Scatter::new(vec![0, 1, 2], vec![6, 10, 2]).name("trace1");
         let mut plot = Plot::new();
@@ -988,6 +981,15 @@ mod tests {
         assert!(!dst.exists());
     }
 
+    #[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 {
+        PORT_COUNTER.fetch_add(1, Ordering::SeqCst)
+    }
+
     #[test]
     #[cfg(feature = "plotly_static")]
     fn save_to_png() {
@@ -1003,8 +1005,8 @@ mod tests {
         let metadata = std::fs::metadata(&dst).expect("Could not retrieve file metadata");
         let file_size = metadata.len();
         assert!(file_size > 0,);
-        // assert!(std::fs::remove_file(&dst).is_ok());
-        // assert!(!dst.exists());
+        assert!(std::fs::remove_file(&dst).is_ok());
+        assert!(!dst.exists());
     }
 
     #[test]

plotly_static/build.rs

@@ -336,13 +336,12 @@ fn main() -> Result<()> {
     } else {
         #[cfg(feature = "chromedriver")]
         {
-            let msg = format!("'webdriver_download' feature disabled. Please install a {CHROMEDRIVER_PATH_ENV} version manually and make the environment variable 'WEBDRIVER_PATH' point to it.");
+            let msg = "'webdriver_download' feature disabled. Please install a 'chromedriver' version manually and make the environment variable 'WEBDRIVER_PATH' point to it.".to_string();
             println!("cargo::warning={msg}");
         }
         #[cfg(feature = "geckodriver")]
         {
-            let msg = format!("'webdriver_download' feature disabled. Please install a {} version manually and make the environment variable 'WEBDRIVER_PATH' point to it.", 
-                GECKODRIVER_PATH_ENV);
+            let msg = format!("'webdriver_download' feature disabled. Please install a 'geckodriver' version manually and make the environment variable 'WEBDRIVER_PATH' point to it.");
             println!("cargo::warning={msg}");
         }
         #[cfg(not(any(feature = "chromedriver", feature = "geckodriver")))]

plotly_static/src/lib.rs

@@ -266,7 +266,7 @@ use std::{println as error, println as warn, println as debug};
 
 use anyhow::{anyhow, Context, Result};
 use base64::{engine::general_purpose, Engine as _};
-use fantoccini::{wd::Capabilities, ClientBuilder};
+use fantoccini::{wd::Capabilities, Client, ClientBuilder};
 #[cfg(not(test))]
 use log::{debug, error, warn};
 use serde::Serialize;
@@ -426,9 +426,8 @@ impl std::fmt::Display for ImageFormat {
     }
 }
 
-// TODO: how to avoid cyclic dependency on the ImageFormat and the Plot data
-// ideally ImageFormat will be defined in a single place and the `data` field
-// would be just a Plot object which is later serialized to JSON
+/// TODO: ideally data would be a Plot object which is later serialized to JSON
+/// but with the current workspace set up, that would be a cyclic dependency.
 #[derive(Serialize)]
 struct PlotData<'a> {
     format: ImageFormat,
@@ -469,10 +468,15 @@ struct PlotData<'a> {
 /// - Browser capabilities: Default Chrome/Firefox headless options
 /// - Automatic WebDriver detection and connection reuse
 pub struct StaticExporterBuilder {
+    /// WebDriver server port (default: 4444)
     webdriver_port: u32,
+    /// WebDriver server base URL (default: "http://localhost")
     webdriver_url: String,
+    /// Auto-spawn WebDriver if not running (default: true)
     spawn_webdriver: bool,
+    /// Use bundled JS libraries instead of CDN (default: false)
     offline_mode: bool,
+    /// Browser command-line flags (e.g., "--headless", "--no-sandbox")
     webdriver_browser_caps: Vec<String>,
 }
 
@@ -675,6 +679,7 @@ impl StaticExporterBuilder {
             offline_mode: self.offline_mode,
             webdriver_browser_caps: self.webdriver_browser_caps.clone(),
             runtime,
+            webdriver_client: None,
         })
     }
 
@@ -737,12 +742,26 @@ impl StaticExporterBuilder {
 /// - Offline mode support
 /// - Automatic WebDriver management
 pub struct StaticExporter {
+    /// WebDriver server port (default: 4444)
     webdriver_port: u32,
+
+    /// WebDriver server base URL (default: "http://localhost")
     webdriver_url: String,
+
+    /// WebDriver process manager for spawning and cleanup
     webdriver: WebDriver,
+
+    /// Use bundled JS libraries instead of CDN
     offline_mode: bool,
+
+    /// Browser command-line flags (e.g., "--headless", "--no-sandbox")
     webdriver_browser_caps: Vec<String>,
+
+    /// Tokio runtime for async operations
     runtime: std::sync::Arc<tokio::runtime::Runtime>,
+
+    /// Cached WebDriver client for session reuse
+    webdriver_client: Option<Client>,
 }
 
 impl Drop for StaticExporter {
@@ -757,6 +776,16 @@ impl Drop for StaticExporter {
     /// - Leaves externally managed WebDriver sessions running
     /// - Logs errors but doesn't panic if cleanup fails
     fn drop(&mut self) {
+        // Close the WebDriver client if it exists
+        if let Some(client) = self.webdriver_client.take() {
+            let runtime = self.runtime.clone();
+            runtime.block_on(async {
+                if let Err(e) = client.close().await {
+                    error!("Failed to close WebDriver client: {e}");
+                }
+            });
+        }
+
         // Stop the WebDriver process
         if let Err(e) = self.webdriver.stop() {
             error!("Failed to stop WebDriver: {e}");
@@ -941,11 +970,20 @@ impl StaticExporter {
         debug!("Use WebDriver and headless browser to export static plot");
         let webdriver_url = format!("{}:{}", self.webdriver_url, self.webdriver_port,);
 
-        let client = ClientBuilder::native()
-            .capabilities(caps)
-            .connect(&webdriver_url)
-            .await
-            .with_context(|| "WebDriver session errror")?;
+        // Reuse existing client or create new one
+        let client = if let Some(ref client) = self.webdriver_client {
+            debug!("Reusing existing WebDriver session");
+            client.clone()
+        } else {
+            debug!("Creating new WebDriver session");
+            let new_client = ClientBuilder::native()
+                .capabilities(caps)
+                .connect(&webdriver_url)
+                .await
+                .with_context(|| "WebDriver session error")?;
+            self.webdriver_client = Some(new_client.clone());
+            new_client
+        };
 
         // URL-encode the HTML
         let data_uri = format!("data:text/html,{}", encode(data_uri));
@@ -978,7 +1016,8 @@ impl StaticExporter {
 
         let data = client.execute_async(js_script, args).await?;
 
-        client.close().await?;
+        // Don't close the client - keep it for reuse
+        // client.close().await?;
 
         let src = data.as_str().ok_or(anyhow!(
             "Failed to execute Plotly.toImage in browser session"
@@ -1282,46 +1321,50 @@ mod tests {
         init();
         let test_plot = create_test_plot();
 
+        // Use a unique port to test actual WebDriver process reuse
+        let test_port = get_unique_port();
+
         // Create first exporter - this should spawn a new WebDriver
         let mut export1 = StaticExporterBuilder::default()
             .spawn_webdriver(true)
-            .webdriver_port(get_unique_port())
+            .webdriver_port(test_port)
             .build()
             .unwrap();
 
         // Export first image
-        let dst1 = PathBuf::from("session_reuse_1.png");
+        let dst1 = PathBuf::from("process_reuse_1.png");
         export1
             .write_fig(dst1.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
             .unwrap();
         assert!(dst1.exists());
         assert!(std::fs::remove_file(dst1.as_path()).is_ok());
 
         // Create second exporter on the same port - this should connect to existing
-        // WebDriver
+        // WebDriver process (but create a new session)
         let mut export2 = StaticExporterBuilder::default()
             .spawn_webdriver(true)
-            .webdriver_port(get_unique_port())
+            .webdriver_port(test_port)
             .build()
             .unwrap();
 
-        // Export second image using the same WebDriver session
-        let dst2 = PathBuf::from("session_reuse_2.png");
+        // Export second image using a new session on the same WebDriver process
+        let dst2 = PathBuf::from("process_reuse_2.png");
         export2
             .write_fig(dst2.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
             .unwrap();
         assert!(dst2.exists());
         assert!(std::fs::remove_file(dst2.as_path()).is_ok());
 
         // Create third exporter on the same port - should also connect to existing
+        // WebDriver process
         let mut export3 = StaticExporterBuilder::default()
             .spawn_webdriver(true)
-            .webdriver_port(get_unique_port())
+            .webdriver_port(test_port)
             .build()
             .unwrap();
 
-        // Export third image
-        let dst3 = PathBuf::from("session_reuse_3.png");
+        // Export third image using another new session on the same WebDriver process
+        let dst3 = PathBuf::from("process_reuse_3.png");
         export3
             .write_fig(dst3.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
             .unwrap();