Skip to content

Commit 09942be

Browse files
andrei-ngandrei-ng
committed
improve webdriver detection/reuse
Signed-off-by: Andrei Gherghescu <[email protected]>
1 parent 9d47885 commit 09942be

File tree

5 files changed

+232
-37
lines changed

5 files changed

+232
-37
lines changed

‎plotly/src/plot.rs‎

Lines changed: 81 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -803,6 +803,7 @@ impl PartialEq for Plot {
803803
#[cfg(test)]
804804
mod tests {
805805
use std::path::PathBuf;
806+
use std::sync::atomic::{AtomicU32, Ordering};
806807

807808
#[cfg(feature = "kaleido")]
808809
use plotly_kaleido::ImageFormat;
@@ -815,6 +816,13 @@ mod tests {
815816
use super::*;
816817
use crate::Scatter;
817818

819+
// Helper to generate unique ports for parallel tests
820+
static PORT_COUNTER: AtomicU32 = AtomicU32::new(4444);
821+
822+
fn get_unique_port() -> u32 {
823+
PORT_COUNTER.fetch_add(1, Ordering::SeqCst)
824+
}
825+
818826
fn create_test_plot() -> Plot {
819827
let trace1 = Scatter::new(vec![0, 1, 2], vec![6, 10, 2]).name("trace1");
820828
let mut plot = Plot::new();
@@ -961,18 +969,25 @@ mod tests {
961969
let dst = PathBuf::from("example.html");
962970
plot.write_html(&dst);
963971
assert!(dst.exists());
964-
// assert!(std::fs::remove_file(&dst).is_ok());
965-
// assert!(!dst.exists());
972+
assert!(std::fs::remove_file(&dst).is_ok());
973+
assert!(!dst.exists());
966974
}
967975

968976
#[test]
969977
#[cfg(feature = "plotly_static")]
970978
fn save_to_png() {
971979
let plot = create_test_plot();
972980
let dst = PathBuf::from("example.png");
973-
plot.write_image(&dst, ImageFormat::PNG, 1024, 680, 1.0)
981+
let mut exporter = plotly_static::StaticExporterBuilder::default()
982+
.webdriver_port(get_unique_port())
983+
.build()
984+
.unwrap();
985+
plot.write_image_with_exporter(&mut exporter, &dst, ImageFormat::PNG, 1024, 680, 1.0)
974986
.unwrap();
975987
assert!(dst.exists());
988+
let metadata = std::fs::metadata(&dst).expect("Could not retrieve file metadata");
989+
let file_size = metadata.len();
990+
assert!(file_size > 0,);
976991
// assert!(std::fs::remove_file(&dst).is_ok());
977992
// assert!(!dst.exists());
978993
}
@@ -982,55 +997,87 @@ mod tests {
982997
fn save_to_jpeg() {
983998
let plot = create_test_plot();
984999
let dst = PathBuf::from("example.jpeg");
985-
plot.write_image(&dst, ImageFormat::JPEG, 1024, 680, 1.0)
1000+
let mut exporter = plotly_static::StaticExporterBuilder::default()
1001+
.webdriver_port(get_unique_port())
1002+
.build()
1003+
.unwrap();
1004+
plot.write_image_with_exporter(&mut exporter, &dst, ImageFormat::JPEG, 1024, 680, 1.0)
9861005
.unwrap();
9871006
assert!(dst.exists());
988-
// assert!(std::fs::remove_file(&dst).is_ok());
989-
// assert!(!dst.exists());
1007+
let metadata = std::fs::metadata(&dst).expect("Could not retrieve file metadata");
1008+
let file_size = metadata.len();
1009+
assert!(file_size > 0,);
1010+
assert!(std::fs::remove_file(&dst).is_ok());
1011+
assert!(!dst.exists());
9901012
}
9911013

9921014
#[test]
9931015
#[cfg(feature = "plotly_static")]
9941016
fn save_to_svg() {
9951017
let plot = create_test_plot();
9961018
let dst = PathBuf::from("example.svg");
997-
plot.write_image(&dst, ImageFormat::SVG, 1024, 680, 1.0)
1019+
let mut exporter = plotly_static::StaticExporterBuilder::default()
1020+
.webdriver_port(get_unique_port())
1021+
.build()
1022+
.unwrap();
1023+
plot.write_image_with_exporter(&mut exporter, &dst, ImageFormat::SVG, 1024, 680, 1.0)
9981024
.unwrap();
9991025
assert!(dst.exists());
1000-
// assert!(std::fs::remove_file(&dst).is_ok());
1001-
// assert!(!dst.exists());
1026+
let metadata = std::fs::metadata(&dst).expect("Could not retrieve file metadata");
1027+
let file_size = metadata.len();
1028+
assert!(file_size > 0,);
1029+
assert!(std::fs::remove_file(&dst).is_ok());
1030+
assert!(!dst.exists());
10021031
}
10031032

10041033
#[test]
10051034
#[cfg(feature = "plotly_static")]
10061035
fn save_to_pdf() {
10071036
let plot = create_test_plot();
10081037
let dst = PathBuf::from("example.pdf");
1009-
plot.write_image(&dst, ImageFormat::PDF, 1024, 680, 1.0)
1038+
let mut exporter = plotly_static::StaticExporterBuilder::default()
1039+
.webdriver_port(get_unique_port())
1040+
.build()
1041+
.unwrap();
1042+
plot.write_image_with_exporter(&mut exporter, &dst, ImageFormat::PDF, 1024, 680, 1.0)
10101043
.unwrap();
10111044
assert!(dst.exists());
1012-
// assert!(std::fs::remove_file(&dst).is_ok());
1013-
// assert!(!dst.exists());
1045+
let metadata = std::fs::metadata(&dst).expect("Could not retrieve file metadata");
1046+
let file_size = metadata.len();
1047+
assert!(file_size > 0,);
1048+
assert!(std::fs::remove_file(&dst).is_ok());
1049+
assert!(!dst.exists());
10141050
}
10151051

10161052
#[test]
10171053
#[cfg(feature = "plotly_static")]
10181054
fn save_to_webp() {
10191055
let plot = create_test_plot();
10201056
let dst = PathBuf::from("example.webp");
1021-
plot.write_image(&dst, ImageFormat::WEBP, 1024, 680, 1.0)
1057+
let mut exporter = plotly_static::StaticExporterBuilder::default()
1058+
.webdriver_port(get_unique_port())
1059+
.build()
1060+
.unwrap();
1061+
plot.write_image_with_exporter(&mut exporter, &dst, ImageFormat::WEBP, 1024, 680, 1.0)
10221062
.unwrap();
10231063
assert!(dst.exists());
1024-
// assert!(std::fs::remove_file(&dst).is_ok());
1025-
// assert!(!dst.exists());
1064+
let metadata = std::fs::metadata(&dst).expect("Could not retrieve file metadata");
1065+
let file_size = metadata.len();
1066+
assert!(file_size > 0,);
1067+
assert!(std::fs::remove_file(&dst).is_ok());
1068+
assert!(!dst.exists());
10261069
}
10271070

10281071
#[test]
10291072
#[cfg(feature = "plotly_static")]
10301073
fn image_to_base64() {
10311074
let plot = create_test_plot();
1075+
let mut exporter = plotly_static::StaticExporterBuilder::default()
1076+
.webdriver_port(get_unique_port())
1077+
.build()
1078+
.unwrap();
10321079

1033-
let image_base64 = plot.to_base64(ImageFormat::PNG, 200, 150, 1.0).unwrap();
1080+
let image_base64 = plot.to_base64_with_exporter(&mut exporter, ImageFormat::PNG, 200, 150, 1.0).unwrap();
10341081

10351082
assert!(!image_base64.is_empty());
10361083

@@ -1048,7 +1095,11 @@ mod tests {
10481095
#[cfg(feature = "plotly_static")]
10491096
fn image_to_svg_string() {
10501097
let plot = create_test_plot();
1051-
let image_svg = plot.to_svg(200, 150, 1.0).unwrap();
1098+
let mut exporter = plotly_static::StaticExporterBuilder::default()
1099+
.webdriver_port(get_unique_port())
1100+
.build()
1101+
.unwrap();
1102+
let image_svg = plot.to_svg_with_exporter(&mut exporter, 200, 150, 1.0).unwrap();
10521103

10531104
assert!(!image_svg.is_empty());
10541105

@@ -1078,13 +1129,22 @@ mod tests {
10781129

10791130
plot.add_trace(surface);
10801131
let dst = PathBuf::from("example.png");
1081-
plot.write_image(&dst, ImageFormat::PNG, 800, 600, 1.0)
1132+
let mut exporter = plotly_static::StaticExporterBuilder::default()
1133+
.webdriver_port(get_unique_port())
1134+
.build()
1135+
.unwrap();
1136+
plot.write_image_with_exporter(&mut exporter, &dst, ImageFormat::PNG, 800, 600, 1.0)
10821137
.unwrap();
10831138
assert!(dst.exists());
1084-
// assert!(std::fs::remove_file(&dst).is_ok());
1085-
// assert!(!dst.exists());
1139+
1140+
let metadata = std::fs::metadata(&dst).expect("Could not retrieve file metadata");
1141+
let file_size = metadata.len();
1142+
assert!(file_size > 0,);
1143+
assert!(std::fs::remove_file(&dst).is_ok());
1144+
assert!(!dst.exists());
1145+
10861146
assert!(!plot
1087-
.to_base64(ImageFormat::PNG, 1024, 680, 1.0)
1147+
.to_base64_with_exporter(&mut exporter, ImageFormat::PNG, 1024, 680, 1.0)
10881148
.unwrap()
10891149
.is_empty());
10901150
}

‎plotly_static/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ fantoccini = "0.21"
2727
tokio = { version = "1", features = ["full"] }
2828
anyhow = "1.0"
2929
urlencoding = "2"
30+
reqwest = { version = "0.11", features = ["blocking"] }
3031

3132
[dev-dependencies]
3233
plotly_static = { path = "." }

‎plotly_static/build.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,8 +116,7 @@ async fn download_with_retry(
116116
last_error = Some(e);
117117
attempts += 1;
118118
if attempts < MAX_DOWNLOAD_RETRIES {
119-
let delay =
120-
Duration::from_secs(INITIAL_RETRY_DELAY * 2u64.pow(attempts - 1));
119+
let delay = Duration::from_secs(INITIAL_RETRY_DELAY * 2u64.pow(attempts - 1));
121120
println!(
122121
"cargo:warning=Download attempt {} failed, retrying in {:?}...",
123122
attempts, delay

‎plotly_static/src/lib.rs‎

Lines changed: 74 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -473,10 +473,14 @@ impl StaticExporterBuilder {
473473
/// .expect("Failed to build StaticExporter");
474474
/// ```
475475
pub fn build(&self) -> Result<StaticExporter> {
476-
let mut wd = WebDriver::new(self.webdriver_port)?;
477-
if self.spawn_webdriver {
478-
wd.spawn_webdriver();
479-
}
476+
let wd = if self.spawn_webdriver {
477+
// Try to connect to existing WebDriver first, spawn new if not available
478+
WebDriver::new_or_connect(self.webdriver_port)?
479+
} else {
480+
// Just create the WebDriver instance without spawning
481+
WebDriver::new(self.webdriver_port)?
482+
};
483+
480484
Ok(StaticExporter {
481485
webdriver_port: self.webdriver_port,
482486
webdriver_url: self.webdriver_url.to_string(),
@@ -854,6 +858,7 @@ impl StaticExporter {
854858
#[cfg(test)]
855859
mod tests {
856860
use std::path::PathBuf;
861+
use std::sync::atomic::{AtomicU32, Ordering};
857862

858863
use env_logger;
859864

@@ -862,6 +867,13 @@ mod tests {
862867
fn init() {
863868
let _ = env_logger::try_init();
864869
}
870+
871+
// Helper to generate unique ports for parallel tests
872+
static PORT_COUNTER: AtomicU32 = AtomicU32::new(4444);
873+
874+
fn get_unique_port() -> u32 {
875+
PORT_COUNTER.fetch_add(1, Ordering::SeqCst)
876+
}
865877

866878
fn create_test_plot() -> serde_json::Value {
867879
serde_json::to_value(serde_json::json!(
@@ -913,7 +925,7 @@ mod tests {
913925

914926
let mut export = StaticExporterBuilder::default()
915927
.spawn_webdriver(true)
916-
.webdriver_port(4444)
928+
.webdriver_port(get_unique_port())
917929
.build()
918930
.unwrap();
919931
let dst = PathBuf::from("example.png");
@@ -933,7 +945,7 @@ mod tests {
933945
let test_plot = create_test_plot();
934946
let mut export = StaticExporterBuilder::default()
935947
.spawn_webdriver(true)
936-
.webdriver_port(4445)
948+
.webdriver_port(get_unique_port())
937949
.build()
938950
.unwrap();
939951
let dst = PathBuf::from("example.jpeg");
@@ -953,7 +965,7 @@ mod tests {
953965
let test_plot = create_test_plot();
954966
let mut export = StaticExporterBuilder::default()
955967
.spawn_webdriver(true)
956-
.webdriver_port(4446)
968+
.webdriver_port(get_unique_port())
957969
.build()
958970
.unwrap();
959971

@@ -984,7 +996,7 @@ mod tests {
984996
let test_plot = create_test_plot();
985997
let mut export = StaticExporterBuilder::default()
986998
.spawn_webdriver(true)
987-
.webdriver_port(4447)
999+
.webdriver_port(get_unique_port())
9881000
.build()
9891001
.unwrap();
9901002
let dst = PathBuf::from("example.svg");
@@ -1004,7 +1016,7 @@ mod tests {
10041016
let test_plot = create_test_plot();
10051017
let mut export = StaticExporterBuilder::default()
10061018
.spawn_webdriver(true)
1007-
.webdriver_port(4449)
1019+
.webdriver_port(get_unique_port())
10081020
.build()
10091021
.unwrap();
10101022
let dst = PathBuf::from("example.webp");
@@ -1024,7 +1036,7 @@ mod tests {
10241036
let test_plot = create_test_plot();
10251037
let mut export = StaticExporterBuilder::default()
10261038
.spawn_webdriver(true)
1027-
.webdriver_port(4450)
1039+
.webdriver_port(get_unique_port())
10281040
.build()
10291041
.unwrap();
10301042
let dst = PathBuf::from("example.pdf");
@@ -1037,6 +1049,58 @@ mod tests {
10371049
assert!(file_size > 0,);
10381050
assert!(std::fs::remove_file(dst.as_path()).is_ok());
10391051
}
1052+
1053+
#[test]
1054+
fn test_webdriver_process_reuse() {
1055+
init();
1056+
let test_plot = create_test_plot();
1057+
1058+
// Create first exporter - this should spawn a new WebDriver
1059+
let mut export1 = StaticExporterBuilder::default()
1060+
.spawn_webdriver(true)
1061+
.webdriver_port(get_unique_port())
1062+
.build()
1063+
.unwrap();
1064+
1065+
// Export first image
1066+
let dst1 = PathBuf::from("session_reuse_1.png");
1067+
export1
1068+
.write_fig(dst1.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
1069+
.unwrap();
1070+
assert!(dst1.exists());
1071+
assert!(std::fs::remove_file(dst1.as_path()).is_ok());
1072+
1073+
// Create second exporter on the same port - this should connect to existing
1074+
// WebDriver
1075+
let mut export2 = StaticExporterBuilder::default()
1076+
.spawn_webdriver(true)
1077+
.webdriver_port(get_unique_port())
1078+
.build()
1079+
.unwrap();
1080+
1081+
// Export second image using the same WebDriver session
1082+
let dst2 = PathBuf::from("session_reuse_2.png");
1083+
export2
1084+
.write_fig(dst2.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
1085+
.unwrap();
1086+
assert!(dst2.exists());
1087+
assert!(std::fs::remove_file(dst2.as_path()).is_ok());
1088+
1089+
// Create third exporter on the same port - should also connect to existing
1090+
let mut export3 = StaticExporterBuilder::default()
1091+
.spawn_webdriver(true)
1092+
.webdriver_port(get_unique_port())
1093+
.build()
1094+
.unwrap();
1095+
1096+
// Export third image
1097+
let dst3 = PathBuf::from("session_reuse_3.png");
1098+
export3
1099+
.write_fig(dst3.as_path(), &test_plot, ImageFormat::PNG, 800, 600, 1.0)
1100+
.unwrap();
1101+
assert!(dst3.exists());
1102+
assert!(std::fs::remove_file(dst3.as_path()).is_ok());
1103+
}
10401104
}
10411105

10421106
#[cfg(feature = "chromedriver")]

0 commit comments

Comments
 (0)