From ca417948ad74e3085d402147d09648adec0d5e4e Mon Sep 17 00:00:00 2001
From: Jon Dufresne <jon.dufresne@gmail.com>
Date: Thu, 18 May 2023 05:34:01 -0700
Subject: [PATCH] Pass window object to mockPrototype (#100)

When mocking the canvas on a non-global window object, the prototype
mocks should apply to the same window object.
---
 __tests__/mock/window.js | 30 ++++++++++++++++++++++++++++++
 src/mock/prototype.js    | 27 ++++++++++++++-------------
 src/window.js            |  2 +-
 3 files changed, 45 insertions(+), 14 deletions(-)
 create mode 100644 __tests__/mock/window.js

diff --git a/__tests__/mock/window.js b/__tests__/mock/window.js
new file mode 100644
index 0000000..b3db593
--- /dev/null
+++ b/__tests__/mock/window.js
@@ -0,0 +1,30 @@
+import { JSDOM } from 'jsdom';
+import mockWindow from '../../src/window';
+
+describe('mockWindow', () => {
+  it('mocks the passed object', () => {
+    const win = new JSDOM().window;
+
+    mockWindow(win);
+
+    expect(win.Path2D).not.toBeNull();
+    expect(win.CanvasGradient).not.toBeNull();
+    expect(win.CanvasPattern).not.toBeNull();
+    expect(win.CanvasRenderingContext2D).not.toBeNull();
+    expect(win.DOMMatrix).not.toBeNull();
+    expect(win.ImageData).not.toBeNull();
+    expect(win.TextMetrics).not.toBeNull();
+    expect(win.ImageBitmap).not.toBeNull();
+    expect(win.createImageBitmap).not.toBeNull();
+
+    expect(
+      jest.isMockFunction(win.HTMLCanvasElement.prototype.getContext)
+    ).toBe(true);
+    expect(jest.isMockFunction(win.HTMLCanvasElement.prototype.toBlob)).toBe(
+      true
+    );
+    expect(jest.isMockFunction(win.HTMLCanvasElement.prototype.toDataURL)).toBe(
+      true
+    );
+  });
+});
diff --git a/src/mock/prototype.js b/src/mock/prototype.js
index 0a733a7..3983654 100644
--- a/src/mock/prototype.js
+++ b/src/mock/prototype.js
@@ -1,4 +1,4 @@
-export default function mockPrototype() {
+export default function mockPrototype(win) {
   /**
    * This weakmap is designed to contain all of the generated canvas contexts. It's keys are the
    * jsdom canvases obtained by using the `this` keyword inside the `#getContext('2d')` function
@@ -29,12 +29,12 @@ export default function mockPrototype() {
     return getContext2D.internal.call(this, type);
   });
 
-  if (!jest.isMockFunction(HTMLCanvasElement.prototype.getContext)) {
-    getContext2D.internal = HTMLCanvasElement.prototype.getContext;
+  if (!jest.isMockFunction(win.HTMLCanvasElement.prototype.getContext)) {
+    getContext2D.internal = win.HTMLCanvasElement.prototype.getContext;
   } else {
-    getContext2D.internal = HTMLCanvasElement.prototype.getContext.internal;
+    getContext2D.internal = win.HTMLCanvasElement.prototype.getContext.internal;
   }
-  HTMLCanvasElement.prototype.getContext = getContext2D;
+  win.HTMLCanvasElement.prototype.getContext = getContext2D;
 
   /**
    * This function technically throws SecurityError at runtime, but it cannot be mocked, because
@@ -73,12 +73,12 @@ export default function mockPrototype() {
     setTimeout(() => callback(blob), 0);
   });
 
-  if (!jest.isMockFunction(HTMLCanvasElement.prototype.toBlob)) {
-    toBlobOverride.internal = HTMLCanvasElement.prototype.toBlob;
+  if (!jest.isMockFunction(win.HTMLCanvasElement.prototype.toBlob)) {
+    toBlobOverride.internal = win.HTMLCanvasElement.prototype.toBlob;
   } else {
-    toBlobOverride.internal = HTMLCanvasElement.prototype.toBlob.internal;
+    toBlobOverride.internal = win.HTMLCanvasElement.prototype.toBlob.internal;
   }
-  HTMLCanvasElement.prototype.toBlob = toBlobOverride;
+  win.HTMLCanvasElement.prototype.toBlob = toBlobOverride;
 
   /**
    * This section creates a dataurl with a validated mime type. This is not actually valid, because
@@ -103,10 +103,11 @@ export default function mockPrototype() {
     return 'data:' + type + ';base64,00';
   });
 
-  if (!jest.isMockFunction(HTMLCanvasElement.prototype.toDataURL)) {
-    toDataURLOverride.internal = HTMLCanvasElement.prototype.toDataURL;
+  if (!jest.isMockFunction(win.HTMLCanvasElement.prototype.toDataURL)) {
+    toDataURLOverride.internal = win.HTMLCanvasElement.prototype.toDataURL;
   } else {
-    toDataURLOverride.internal = HTMLCanvasElement.prototype.toDataURL.internal;
+    toDataURLOverride.internal =
+      win.HTMLCanvasElement.prototype.toDataURL.internal;
   }
-  HTMLCanvasElement.prototype.toDataURL = toDataURLOverride;
+  win.HTMLCanvasElement.prototype.toDataURL = toDataURLOverride;
 }
diff --git a/src/window.js b/src/window.js
index 74d8e72..9c79351 100644
--- a/src/window.js
+++ b/src/window.js
@@ -55,7 +55,7 @@ export default (win) => {
   if (!win.ImageBitmap) win.ImageBitmap = ImageBitmap;
   if (!win.createImageBitmap) win.createImageBitmap = createImageBitmap;
 
-  mockPrototype();
+  mockPrototype(win);
 
   return win;
 };