♻️ Refactor code to make issues easier to fix (#155)

frinyvonnick · Feb 8, 2022 · d28d88e · d28d88e
1 parent c42b6f0
commit d28d88e
Show file tree

Hide file tree

Showing 7 changed files with 379 additions and 138 deletions.
diff --git a/src/main.spec.ts b/src/main.spec.ts
@@ -5,14 +5,44 @@ import { createWorker } from "tesseract.js";
 import { nodeHtmlToImage } from "./main";
 
 describe("node-html-to-image", () => {
+  let mockExit;
+  let mockConsoleErr;
+  const originalConsoleError = console.error;
   beforeEach(() => {
     rimraf.sync("./generated");
     mkdirSync("./generated");
+    mockExit = jest.spyOn(process, "exit").mockImplementation((number) => {
+      throw new Error("process.exit: " + number);
+    });
+    mockConsoleErr = jest
+      .spyOn(console, "error")
+      .mockImplementation((value) => originalConsoleError(value));
+  });
+
+  afterEach(() => {
+    mockExit.mockRestore();
+    mockConsoleErr.mockRestore();
   });
 
   afterAll(() => {
     rimraf.sync("./generated");
   });
+  describe("error", () => {
+    it("should stop the program properly", async () => {
+      /* eslint-disable @typescript-eslint/ban-ts-comment */
+      await expect(async () => {
+        await nodeHtmlToImage({
+          html: "<html></html>",
+          type: "jpeg",
+          // @ts-ignore
+          quality: "wrong value",
+        });
+      }).rejects.toThrow();
+
+      expect(mockExit).toHaveBeenCalledWith(1);
+      /* eslint-enable @typescript-eslint/ban-ts-comment */
+    });
+  });
 
   describe("single image", () => {
     it("should generate output file", async () => {
@@ -33,18 +63,16 @@ describe("node-html-to-image", () => {
     });
 
     it("should throw an error if html is not provided", async () => {
-      let error;
-      try {
-        // We want to test this for JS users
+      await expect(async () => {
         // eslint-disable-next-line @typescript-eslint/ban-ts-comment
         // @ts-ignore
         await nodeHtmlToImage({
           output: "./generated/image.png",
         });
-      } catch (e) {
-        error = e;
-      }
-      expect(error.message).toEqual(expect.stringContaining("html"));
+      }).rejects.toThrow();
+      expect(mockConsoleErr).toHaveBeenCalledWith(
+        new Error("You must provide an html property.")
+      );
     });
 
     it("should generate an jpeg image", async () => {

diff --git a/src/main.ts b/src/main.ts
@@ -1,61 +1,65 @@
 import { Cluster } from "puppeteer-cluster";
 
-import { Options } from "./types";
+import { Options, ScreenshotParams } from "./types";
 import { makeScreenshot } from "./screenshot";
+import { Screenshot } from "./models/Screenshot";
 
 export async function nodeHtmlToImage(options: Options) {
   const {
     html,
+    encoding,
+    transparent,
     content,
     output,
-    selector = "body",
+    selector,
+    type,
+    quality,
     puppeteerArgs = {},
   } = options;
 
-  if (!html) {
-    throw Error("You must provide an html property.");
-  }
-
-  const cluster: Cluster<{
-    content: object;
-    output: string;
-    selector: string;
-  }> = await Cluster.launch({
+  const cluster: Cluster<ScreenshotParams> = await Cluster.launch({
     concurrency: Cluster.CONCURRENCY_CONTEXT,
     maxConcurrency: 2,
     puppeteerOptions: { ...puppeteerArgs, headless: true },
   });
 
-  const buffers: Array<Buffer | string> = [];
-
-  await cluster.task(async ({ page, data: { content, output, selector } }) => {
-    const buffer = await makeScreenshot(page, {
-      ...options,
-      content,
-      output,
-      selector,
-    });
-    buffers.push(buffer);
-  });
-
-  cluster.on("taskerror", (err) => {
-    throw err;
-  });
-
   const shouldBatch = Array.isArray(content);
   const contents = shouldBatch ? content : [{ ...content, output, selector }];
 
-  contents.forEach((content) => {
-    const { output, selector: contentSelector, ...pageContent } = content;
-    cluster.queue({
-      output,
-      content: pageContent,
-      selector: contentSelector ? contentSelector : selector,
-    });
-  });
+  try {
+    const screenshots: Array<Screenshot> = await Promise.all(
+      contents.map((content) => {
+        const { output, selector: contentSelector, ...pageContent } = content;
+        return cluster.execute(
+          {
+            html,
+            encoding,
+            transparent,
+            output,
+            content: pageContent,
+            selector: contentSelector ? contentSelector : selector,
+            type,
+            quality,
+          },
+          async ({ page, data }) => {
+            const screenshot = await makeScreenshot(page, {
+              ...options,
+              screenshot: new Screenshot(data),
+            });
+            return screenshot;
+          }
+        );
+      })
+    );
+    await cluster.idle();
+    await cluster.close();
 
-  await cluster.idle();
-  await cluster.close();
-
-  return shouldBatch ? buffers : buffers[0];
+    return shouldBatch
+      ? screenshots.map(({ buffer }) => buffer)
+      : screenshots[0].buffer;
+  } catch (err) {
+    console.error(err);
+    await cluster.close();
+    process.exit(1);
+  }
 }
diff --git a/src/models/Screenshot.spec.ts b/src/models/Screenshot.spec.ts
@@ -0,0 +1,161 @@
+import { describe } from "jest-circus";
+import { Screenshot } from "./Screenshot";
+
+describe("Screenshot", () => {
+  const html = "<html><body>Hello world!</body></html>";
+
+  it("should create a minimal Screenshot", () => {
+    const screenshot = new Screenshot({
+      html,
+    });
+
+    expect(screenshot.html).toEqual(html);
+  });
+
+  it("should store buffer", () => {
+    const screenshot = new Screenshot({
+      html,
+    });
+
+    const expectedBuffer = Buffer.alloc(1);
+
+    screenshot.setBuffer(expectedBuffer);
+
+    expect(screenshot.buffer).toEqual(expectedBuffer);
+  });
+
+  it("should set html", () => {
+    const screenshot = new Screenshot({
+      html,
+    });
+
+    const expectedHtml = "<div>Hello</div>";
+
+    screenshot.setHTML(expectedHtml);
+
+    expect(screenshot.html).toEqual(expectedHtml);
+  });
+
+  it("should throw an error when setting html with nothing", () => {
+    const screenshot = new Screenshot({
+      html,
+    });
+
+    expect(() => {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      screenshot.setHTML();
+    }).toThrow("You must provide an html property.");
+
+    expect(screenshot.html).toEqual(html);
+  });
+
+  it("should throw an error when setting html with null", () => {
+    const screenshot = new Screenshot({
+      html,
+    });
+
+    expect(() => {
+      screenshot.setHTML(null);
+    }).toThrow("You must provide an html property.");
+
+    expect(screenshot.html).toEqual(html);
+  });
+
+  it("should throw an error when setting html with empty string", () => {
+    const screenshot = new Screenshot({
+      html,
+    });
+
+    expect(() => {
+      screenshot.setHTML("");
+    }).toThrow("You must provide an html property.");
+
+    expect(screenshot.html).toEqual(html);
+  });
+
+  it("should create a Screenshot with a few attributes", () => {
+    const attributes = {
+      html,
+      output: "something",
+      content: { hello: "Salut" },
+      encoding: "base64",
+    } as const;
+    const screenshot = new Screenshot(attributes);
+
+    expect(screenshot).toEqual({
+      ...attributes,
+      transparent: false,
+      selector: "body",
+      type: "png",
+    });
+  });
+
+  it("should create a Screenshot with different value from defaults", () => {
+    const attributes = {
+      html,
+      transparent: true,
+      selector: "something",
+      type: "jpeg",
+    } as const;
+    const screenshot = new Screenshot(attributes);
+
+    expect(screenshot).toEqual(expect.objectContaining(attributes));
+  });
+
+  it("should not set quality if type is different from jpeg", () => {
+    const screenshot = new Screenshot({
+      html,
+      quality: 3,
+    });
+
+    expect(screenshot.quality).toEqual(undefined);
+  });
+
+  it("should set quality if type is jpeg", () => {
+    const screenshot = new Screenshot({
+      html,
+      type: "jpeg",
+      quality: 3,
+    });
+
+    expect(screenshot.quality).toEqual(3);
+  });
+
+  it("should set quality by default if type is jpeg", () => {
+    const screenshot = new Screenshot({
+      html,
+      type: "jpeg",
+    });
+
+    expect(screenshot.quality).toEqual(80);
+  });
+
+  it("should throw an Error if no params are passed", () => {
+    expect(() => {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      new Screenshot();
+    }).toThrow("You must provide an html property.");
+  });
+
+  it("should throw an Error if html is missing", () => {
+    expect(() => {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      new Screenshot({});
+    }).toThrow("You must provide an html property.");
+  });
+
+  it("should throw an Error if html is null", () => {
+    expect(() => {
+      new Screenshot({ html: null });
+    }).toThrow("You must provide an html property.");
+  });
+
+  it("should throw an Error if html is empty", () => {
+    expect(() => {
+      new Screenshot({ html: "" });
+    }).toThrow("You must provide an html property.");
+  });
+});
diff --git a/src/models/Screenshot.ts b/src/models/Screenshot.ts
@@ -0,0 +1,50 @@
+import { ImageType, Encoding, Content, ScreenshotParams } from "../types";
+
+export class Screenshot {
+  output: string;
+  content: Content;
+  selector: string;
+  html: string;
+  quality?: number;
+  buffer?: Buffer | string;
+  type?: ImageType;
+  encoding?: Encoding;
+  transparent?: boolean;
+
+  constructor(params: ScreenshotParams) {
+    if (!params || !params.html) {
+      throw Error("You must provide an html property.");
+    }
+
+    const {
+      html,
+      encoding,
+      transparent = false,
+      output,
+      content,
+      selector = "body",
+      quality = 80,
+      type = "png",
+    } = params;
+
+    this.html = html;
+    this.encoding = encoding;
+    this.transparent = transparent;
+    this.type = type;
+    this.output = output;
+    this.content = content;
+    this.selector = selector;
+    this.quality = type === "jpeg" ? quality : undefined;
+  }
+
+  setHTML(html: string) {
+    if (!html) {
+      throw Error("You must provide an html property.");
+    }
+    this.html = html;
+  }
+
+  setBuffer(buffer: Buffer | string) {
+    this.buffer = buffer;
+  }
+}