From 8312b75d413d74b6085c0ab2134210982232ed39 Mon Sep 17 00:00:00 2001
From: Timothy Sullivan <tsullivan@elastic.co>
Date: Tue, 7 Dec 2021 11:49:58 -0700
Subject: [PATCH] improve types and add test

---
 .../screenshotting/server/plugin.test.ts      | 67 +++++++++++++++++++
 .../plugins/screenshotting/server/plugin.ts   | 32 +++++----
 2 files changed, 87 insertions(+), 12 deletions(-)
 create mode 100644 x-pack/plugins/screenshotting/server/plugin.test.ts

diff --git a/x-pack/plugins/screenshotting/server/plugin.test.ts b/x-pack/plugins/screenshotting/server/plugin.test.ts
new file mode 100644
index 0000000000000..8bc476380a9ce
--- /dev/null
+++ b/x-pack/plugins/screenshotting/server/plugin.test.ts
@@ -0,0 +1,67 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+jest.mock('./browsers/install');
+
+import type { ScreenshotModePluginSetup } from 'src/plugins/screenshot_mode/server';
+import { CoreSetup, CoreStart, PluginInitializerContext } from 'kibana/server';
+import { coreMock } from 'src/core/server/mocks';
+import { ScreenshottingPlugin } from './plugin';
+import { install } from './browsers/install';
+
+let initContext: PluginInitializerContext;
+let coreSetup: CoreSetup;
+let coreStart: CoreStart;
+let setupDeps: {
+  screenshotMode: ScreenshotModePluginSetup;
+};
+
+beforeEach(() => {
+  const configSchema = {
+    browser: { chromium: { disableSandbox: false } },
+  };
+  initContext = coreMock.createPluginInitializerContext(configSchema);
+  coreSetup = coreMock.createSetup({});
+  coreStart = coreMock.createStart();
+  setupDeps = {
+    screenshotMode: {} as ScreenshotModePluginSetup,
+  };
+});
+
+test('sets up and starts properly', async () => {
+  const plugin = new ScreenshottingPlugin(initContext);
+  const setupContract = plugin.setup(coreSetup, setupDeps);
+  expect(setupContract).toMatchInlineSnapshot(`Object {}`);
+
+  await coreSetup.getStartServices();
+
+  const startContract = plugin.start(coreStart);
+  expect(startContract).toMatchInlineSnapshot(`
+    Object {
+      "diagnose": [Function],
+      "getScreenshots": [Function],
+    }
+  `);
+});
+
+test('handles setup issues', async () => {
+  const plugin = new ScreenshottingPlugin(initContext);
+  (install as jest.Mock).mockRejectedValue(`Unsupported platform!!!`);
+
+  const setupContract = plugin.setup(coreSetup, setupDeps);
+  expect(setupContract).toMatchInlineSnapshot(`Object {}`);
+
+  await coreSetup.getStartServices();
+
+  const startContract = plugin.start(coreStart);
+  expect(startContract).toMatchInlineSnapshot(`
+    Object {
+      "diagnose": [Function],
+      "getScreenshots": [Function],
+    }
+  `);
+});
diff --git a/x-pack/plugins/screenshotting/server/plugin.ts b/x-pack/plugins/screenshotting/server/plugin.ts
index f1b7c21e460ef..dfe41bec94a91 100755
--- a/x-pack/plugins/screenshotting/server/plugin.ts
+++ b/x-pack/plugins/screenshotting/server/plugin.ts
@@ -45,7 +45,8 @@ export class ScreenshottingPlugin implements Plugin<void, ScreenshottingStart, S
   private config: ConfigType;
   private logger: Logger;
   private screenshotMode!: ScreenshotModePluginSetup;
-  private browserDriverFactory!: Promise<HeadlessChromiumDriverFactory>;
+  private browserDriverFactory!: Promise<HeadlessChromiumDriverFactory | undefined>;
+  private error?: Error;
 
   constructor(context: PluginInitializerContext<ConfigType>) {
     this.logger = context.logger.get();
@@ -65,25 +66,32 @@ export class ScreenshottingPlugin implements Plugin<void, ScreenshottingStart, S
 
         return new HeadlessChromiumDriverFactory(this.screenshotMode, config, logger, binaryPath);
       } catch (error) {
+        this.error = error;
         this.logger.error('Error in screenshotting setup, it may not function properly.');
         this.logger.error(error);
       }
-    })() as Promise<HeadlessChromiumDriverFactory>;
+    })();
 
     return {};
   }
 
   start({}: CoreStart): ScreenshottingStart {
-    return (
-      this.browserDriverFactory && {
-        diagnose: () =>
-          from(this.browserDriverFactory).pipe(switchMap((factory) => factory.diagnose())),
-        getScreenshots: (options) =>
-          from(this.browserDriverFactory).pipe(
-            switchMap((factory) => getScreenshots(factory, this.logger.get('screenshot'), options))
-          ),
-      }
-    );
+    return {
+      diagnose: () => {
+        if (this.error) {
+          throw this.error;
+        }
+        return from(this.browserDriverFactory).pipe(switchMap((factory) => factory!.diagnose()));
+      },
+      getScreenshots: (options) => {
+        if (this.error) {
+          throw this.error;
+        }
+        return from(this.browserDriverFactory).pipe(
+          switchMap((factory) => getScreenshots(factory!, this.logger.get('screenshot'), options))
+        );
+      },
+    };
   }
 
   stop() {}