use std::path::Path;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
use serde_json::{Map, Value, json};
use crate::error::{Result, ZendriverError};
use crate::tab::Tab;
#[derive(Debug)]
pub struct PdfBuilder<'tab> {
tab: &'tab Tab,
landscape: Option<bool>,
print_background: Option<bool>,
scale: Option<f64>,
paper_width: Option<f64>,
paper_height: Option<f64>,
margin_top: Option<f64>,
margin_bottom: Option<f64>,
margin_left: Option<f64>,
margin_right: Option<f64>,
page_ranges: Option<String>,
prefer_css_page_size: Option<bool>,
}
impl<'tab> PdfBuilder<'tab> {
#[must_use]
pub fn new(tab: &'tab Tab) -> Self {
Self {
tab,
landscape: None,
print_background: None,
scale: None,
paper_width: None,
paper_height: None,
margin_top: None,
margin_bottom: None,
margin_left: None,
margin_right: None,
page_ranges: None,
prefer_css_page_size: None,
}
}
#[must_use]
pub fn landscape(mut self, on: bool) -> Self {
self.landscape = Some(on);
self
}
#[must_use]
pub fn print_background(mut self, on: bool) -> Self {
self.print_background = Some(on);
self
}
#[must_use]
pub fn scale(mut self, scale: f64) -> Self {
self.scale = Some(scale);
self
}
#[must_use]
pub fn paper_width(mut self, inches: f64) -> Self {
self.paper_width = Some(inches);
self
}
#[must_use]
pub fn paper_height(mut self, inches: f64) -> Self {
self.paper_height = Some(inches);
self
}
#[must_use]
pub fn margin_top(mut self, inches: f64) -> Self {
self.margin_top = Some(inches);
self
}
#[must_use]
pub fn margin_bottom(mut self, inches: f64) -> Self {
self.margin_bottom = Some(inches);
self
}
#[must_use]
pub fn margin_left(mut self, inches: f64) -> Self {
self.margin_left = Some(inches);
self
}
#[must_use]
pub fn margin_right(mut self, inches: f64) -> Self {
self.margin_right = Some(inches);
self
}
#[must_use]
pub fn page_ranges(mut self, ranges: impl Into<String>) -> Self {
self.page_ranges = Some(ranges.into());
self
}
#[must_use]
pub fn prefer_css_page_size(mut self, on: bool) -> Self {
self.prefer_css_page_size = Some(on);
self
}
pub async fn bytes(self) -> Result<Vec<u8>> {
let mut params = Map::new();
if let Some(v) = self.landscape {
params.insert("landscape".to_string(), Value::Bool(v));
}
if let Some(v) = self.print_background {
params.insert("printBackground".to_string(), Value::Bool(v));
}
if let Some(v) = self.scale {
params.insert("scale".to_string(), json!(v));
}
if let Some(v) = self.paper_width {
params.insert("paperWidth".to_string(), json!(v));
}
if let Some(v) = self.paper_height {
params.insert("paperHeight".to_string(), json!(v));
}
if let Some(v) = self.margin_top {
params.insert("marginTop".to_string(), json!(v));
}
if let Some(v) = self.margin_bottom {
params.insert("marginBottom".to_string(), json!(v));
}
if let Some(v) = self.margin_left {
params.insert("marginLeft".to_string(), json!(v));
}
if let Some(v) = self.margin_right {
params.insert("marginRight".to_string(), json!(v));
}
if let Some(v) = self.page_ranges {
params.insert("pageRanges".to_string(), Value::String(v));
}
if let Some(v) = self.prefer_css_page_size {
params.insert("preferCSSPageSize".to_string(), Value::Bool(v));
}
let res = self
.tab
.call("Page.printToPDF", Value::Object(params))
.await?;
let data = res
.get("data")
.and_then(|v| v.as_str())
.ok_or_else(|| ZendriverError::Navigation("Page.printToPDF returned no data".into()))?;
BASE64
.decode(data)
.map_err(|e| ZendriverError::Navigation(format!("invalid base64 in pdf: {e}")))
}
pub async fn save(self, path: impl AsRef<Path>) -> Result<()> {
let bytes = self.bytes().await?;
tokio::fs::write(path, bytes).await?;
Ok(())
}
}
impl Tab {
#[must_use]
pub fn pdf_builder(&self) -> PdfBuilder<'_> {
PdfBuilder::new(self)
}
pub async fn print_to_pdf(&self, path: impl AsRef<Path>) -> Result<()> {
self.pdf_builder()
.paper_width(8.27)
.paper_height(11.69)
.save(path)
.await
}
pub async fn snapshot_mhtml(&self) -> Result<String> {
let res = self
.call("Page.captureSnapshot", json!({ "format": "mhtml" }))
.await?;
let data = res.get("data").and_then(|v| v.as_str()).ok_or_else(|| {
ZendriverError::Navigation("Page.captureSnapshot returned no data".into())
})?;
Ok(data.to_string())
}
pub async fn save_snapshot(&self, path: impl AsRef<Path>) -> Result<()> {
let mhtml = self.snapshot_mhtml().await?;
tokio::fs::write(path, mhtml).await?;
Ok(())
}
}
#[cfg(test)]
#[allow(clippy::panic, clippy::unwrap_used)]
mod tests {
use super::*;
use serde_json::json;
use zendriver_transport::SessionHandle;
use zendriver_transport::testing::MockConnection;
#[tokio::test]
async fn pdf_builder_landscape_dispatches_print_to_pdf() {
let (mut mock, conn) = MockConnection::pair();
let sess = SessionHandle::new(conn.clone(), "S1");
let tab = Tab::new_for_test(sess);
let fut = tokio::spawn({
let t = tab.clone();
async move { PdfBuilder::new(&t).landscape(true).bytes().await }
});
let id = mock.expect_cmd("Page.printToPDF").await;
let sent = mock.last_sent();
assert_eq!(sent["params"]["landscape"], true);
mock.reply(id, json!({ "data": "JVBERg==" })).await;
let bytes = fut.await.unwrap().unwrap();
assert_eq!(bytes, b"%PDF");
conn.shutdown();
}
#[tokio::test]
async fn print_to_pdf_saves_file() {
let (mut mock, conn) = MockConnection::pair();
let sess = SessionHandle::new(conn.clone(), "S1");
let tab = Tab::new_for_test(sess);
let dir = std::env::temp_dir();
let path = dir.join(format!("zendriver_pdf_test_{}.pdf", std::process::id()));
let path_for_task = path.clone();
let fut = tokio::spawn({
let t = tab.clone();
async move { t.print_to_pdf(&path_for_task).await }
});
let id = mock.expect_cmd("Page.printToPDF").await;
let sent = mock.last_sent();
assert_eq!(sent["params"]["paperWidth"], 8.27);
assert_eq!(sent["params"]["paperHeight"], 11.69);
mock.reply(id, json!({ "data": "JVBERg==" })).await;
fut.await.unwrap().unwrap();
let written = std::fs::read(&path).unwrap();
assert_eq!(written, b"%PDF");
let _ = std::fs::remove_file(&path);
conn.shutdown();
}
#[tokio::test]
async fn save_snapshot_dispatches_capture_snapshot_mhtml() {
let (mut mock, conn) = MockConnection::pair();
let sess = SessionHandle::new(conn.clone(), "S1");
let tab = Tab::new_for_test(sess);
let dir = std::env::temp_dir();
let path = dir.join(format!("zendriver_mhtml_test_{}.mhtml", std::process::id()));
let path_for_task = path.clone();
let fut = tokio::spawn({
let t = tab.clone();
async move { t.save_snapshot(&path_for_task).await }
});
let id = mock.expect_cmd("Page.captureSnapshot").await;
let sent = mock.last_sent();
assert_eq!(sent["params"]["format"], "mhtml");
mock.reply(
id,
json!({ "data": "From: <Saved by Chrome>\r\nMIME-Version: 1.0\r\n" }),
)
.await;
fut.await.unwrap().unwrap();
let written = std::fs::read_to_string(&path).unwrap();
assert!(written.contains("MIME-Version"));
let _ = std::fs::remove_file(&path);
conn.shutdown();
}
}