From 824b1aeb711d58bc312a2ebb2d735329c2a59733 Mon Sep 17 00:00:00 2001
From: "trop[bot]" <37223003+trop[bot]@users.noreply.github.com>
Date: Fri, 23 Aug 2024 17:16:22 -0500
Subject: [PATCH] test: add -pdf-ready-to-print event to WebContents for
 testing (#43465)

Co-authored-by: trop[bot] <37223003+trop[bot]@users.noreply.github.com>
Co-authored-by: John Kleinschmidt <jkleinsc@electronjs.org>
---
 BUILD.gn                                      |  1 +
 .../browser/api/electron_api_web_contents.cc  |  4 ++++
 shell/browser/api/electron_api_web_contents.h |  2 ++
 .../electron_pdf_document_helper_client.cc    | 23 +++++++++++++++++++
 .../electron_pdf_document_helper_client.h     |  3 ++-
 spec/api-web-contents-spec.ts                 |  8 ++-----
 6 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/BUILD.gn b/BUILD.gn
index f843ee6c573..30eb55b0877 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -733,6 +733,7 @@ source_set("electron_lib") {
       "//components/pdf/common:util",
       "//components/pdf/renderer",
       "//pdf",
+      "//pdf:content_restriction",
     ]
     sources += [
       "shell/browser/electron_pdf_document_helper_client.cc",
diff --git a/shell/browser/api/electron_api_web_contents.cc b/shell/browser/api/electron_api_web_contents.cc
index c64a7f31b56..b5c3299b414 100644
--- a/shell/browser/api/electron_api_web_contents.cc
+++ b/shell/browser/api/electron_api_web_contents.cc
@@ -3754,6 +3754,10 @@ void WebContents::SetBackgroundColor(std::optional<SkColor> maybe_color) {
   }
 }
 
+void WebContents::PDFReadyToPrint() {
+  Emit("-pdf-ready-to-print");
+}
+
 void WebContents::OnInputEvent(const blink::WebInputEvent& event) {
   Emit("input-event", event);
 }
diff --git a/shell/browser/api/electron_api_web_contents.h b/shell/browser/api/electron_api_web_contents.h
index 90fd58fb68d..2998d820c37 100644
--- a/shell/browser/api/electron_api_web_contents.h
+++ b/shell/browser/api/electron_api_web_contents.h
@@ -489,6 +489,8 @@ class WebContents : public ExclusiveAccessContext,
 
   void SetBackgroundColor(std::optional<SkColor> color);
 
+  void PDFReadyToPrint();
+
   SkRegion* draggable_region() {
     return force_non_draggable_ ? nullptr : draggable_region_.get();
   }
diff --git a/shell/browser/electron_pdf_document_helper_client.cc b/shell/browser/electron_pdf_document_helper_client.cc
index 03db51634af..6902f4c22f5 100644
--- a/shell/browser/electron_pdf_document_helper_client.cc
+++ b/shell/browser/electron_pdf_document_helper_client.cc
@@ -5,6 +5,29 @@
 #include "shell/browser/electron_pdf_document_helper_client.h"
 
 #include "content/public/browser/web_contents.h"
+#include "pdf/content_restriction.h"
+#include "shell/browser/api/electron_api_web_contents.h"
 
 ElectronPDFDocumentHelperClient::ElectronPDFDocumentHelperClient() = default;
 ElectronPDFDocumentHelperClient::~ElectronPDFDocumentHelperClient() = default;
+
+void ElectronPDFDocumentHelperClient::UpdateContentRestrictions(
+    content::RenderFrameHost* render_frame_host,
+    int content_restrictions) {
+  // UpdateContentRestrictions potentially gets called twice from
+  // pdf/pdf_view_web_plugin.cc.  The first time it is potentially called is
+  // when loading starts and it is called with a restriction on printing.  The
+  // second time it is called is when loading is finished and if printing is
+  // allowed there won't be a printing restriction passed, so we can use this
+  // second call to notify that the pdf document is ready to print.
+  if (!(content_restrictions & chrome_pdf::kContentRestrictionPrint)) {
+    content::WebContents* web_contents =
+        content::WebContents::FromRenderFrameHost(render_frame_host);
+    electron::api::WebContents* api_web_contents =
+        electron::api::WebContents::From(
+            web_contents->GetOutermostWebContents());
+    if (api_web_contents) {
+      api_web_contents->PDFReadyToPrint();
+    }
+  }
+}
diff --git a/shell/browser/electron_pdf_document_helper_client.h b/shell/browser/electron_pdf_document_helper_client.h
index b8ed620c7b2..73265544de3 100644
--- a/shell/browser/electron_pdf_document_helper_client.h
+++ b/shell/browser/electron_pdf_document_helper_client.h
@@ -17,8 +17,9 @@ class ElectronPDFDocumentHelperClient : public pdf::PDFDocumentHelperClient {
 
  private:
   // pdf::PDFDocumentHelperClient
+
   void UpdateContentRestrictions(content::RenderFrameHost* render_frame_host,
-                                 int content_restrictions) override {}
+                                 int content_restrictions) override;
   void OnPDFHasUnsupportedFeature(content::WebContents* contents) override {}
   void OnSaveURL(content::WebContents* contents) override {}
   void SetPluginCanSave(content::RenderFrameHost* render_frame_host,
diff --git a/spec/api-web-contents-spec.ts b/spec/api-web-contents-spec.ts
index 5e699880b6f..32cf1dff03e 100644
--- a/spec/api-web-contents-spec.ts
+++ b/spec/api-web-contents-spec.ts
@@ -2450,13 +2450,9 @@ describe('webContents module', () => {
 
     it('from an existing pdf document', async () => {
       const pdfPath = path.join(fixturesPath, 'cat.pdf');
+      const readyToPrint = once(w.webContents, '-pdf-ready-to-print');
       await w.loadFile(pdfPath);
-
-      // TODO(codebytere): the PDF plugin is not always ready immediately
-      // after the document is loaded, so we need to wait for it to be ready.
-      // We should find a better way to do this.
-      await setTimeout(3000);
-
+      await readyToPrint;
       const data = await w.webContents.printToPDF({});
       const doc = await pdfjs.getDocument(data).promise;
       expect(doc.numPages).to.equal(2);