From 3f5b5492831622e68607ef29d4c11561a1abe1e4 Mon Sep 17 00:00:00 2001
From: Daniel Izdebski <spidiritsu@gmail.com>
Date: Wed, 25 Oct 2023 02:26:17 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(core):=20Add=20`waitForExtensi?=
 =?UTF-8?q?onOnLoadPage`=20function=20(#954)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/utils/waitForExtensionOnLoadPage.ts   | 13 +++++
 .../utils/waitForExtensionOnLoadPage.test.ts  | 53 +++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 packages/core/src/utils/waitForExtensionOnLoadPage.ts
 create mode 100644 packages/core/test/utils/waitForExtensionOnLoadPage.test.ts

diff --git a/packages/core/src/utils/waitForExtensionOnLoadPage.ts b/packages/core/src/utils/waitForExtensionOnLoadPage.ts
new file mode 100644
index 000000000..aee02e627
--- /dev/null
+++ b/packages/core/src/utils/waitForExtensionOnLoadPage.ts
@@ -0,0 +1,13 @@
+import { type BrowserContext, errors as playwrightErrors } from 'playwright-core'
+
+export async function waitForExtensionOnLoadPage(context: BrowserContext) {
+  try {
+    return await context.waitForEvent('page', { timeout: 5000 }) // TODO: Extract & Make this timeout configurable.
+  } catch (e) {
+    if (e instanceof playwrightErrors.TimeoutError) {
+      throw new Error('[WaitForExtensionOnLoadPage] Extension did not load in time!')
+    }
+
+    throw e
+  }
+}
diff --git a/packages/core/test/utils/waitForExtensionOnLoadPage.test.ts b/packages/core/test/utils/waitForExtensionOnLoadPage.test.ts
new file mode 100644
index 000000000..2f9b4e01f
--- /dev/null
+++ b/packages/core/test/utils/waitForExtensionOnLoadPage.test.ts
@@ -0,0 +1,53 @@
+import { errors as playwrightErrors } from 'playwright-core'
+import { afterAll, afterEach, describe, expect, it, vi } from 'vitest'
+import { waitForExtensionOnLoadPage } from '../../src/utils/waitForExtensionOnLoadPage'
+
+const createContext = vi.fn().mockReturnValue({
+  waitForEvent: vi.fn().mockResolvedValue(undefined)
+})
+
+const createTimeoutContext = vi.fn().mockReturnValue({
+  waitForEvent: vi.fn().mockImplementation(() => {
+    throw new playwrightErrors.TimeoutError('Timeout :)')
+  })
+})
+
+const createThrowContext = vi.fn().mockReturnValue({
+  waitForEvent: vi.fn().mockImplementation(() => {
+    throw new Error('Unknown Quack!')
+  })
+})
+
+describe('waitForExtensionOnLoadPage', () => {
+  afterAll(() => {
+    vi.resetAllMocks()
+  })
+
+  afterEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('calls context.waitForEvent with correct arguments', async () => {
+    const context = createContext()
+    const waitForEventSpy = vi.spyOn(context, 'waitForEvent')
+
+    await waitForExtensionOnLoadPage(context)
+
+    expect(waitForEventSpy).toHaveBeenCalledOnce()
+    expect(waitForEventSpy).toHaveBeenCalledWith('page', { timeout: 5000 })
+  })
+
+  it('throws with custom error if waitForEvent throws due to timeout', async () => {
+    const context = createTimeoutContext()
+
+    await expect(waitForExtensionOnLoadPage(context)).rejects.toThrowError(
+      '[WaitForExtensionOnLoadPage] Extension did not load in time!'
+    )
+  })
+
+  it('throws if unknown error', async () => {
+    const context = createThrowContext()
+
+    await expect(waitForExtensionOnLoadPage(context)).rejects.toThrowError('Unknown Quack!')
+  })
+})