From 967778ca1881e219f2dd222f003ae05eb71bbe6a Mon Sep 17 00:00:00 2001
From: Gerardo Rodriguez <mr.gerardo.rodriguez@gmail.com>
Date: Tue, 31 Aug 2021 16:11:28 -0700
Subject: [PATCH 1/3] Add a test

---
 tests/utils/runJS.test.tsx | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tests/utils/runJS.test.tsx b/tests/utils/runJS.test.tsx
index 2cd2964f..b5ae1bad 100644
--- a/tests/utils/runJS.test.tsx
+++ b/tests/utils/runJS.test.tsx
@@ -35,6 +35,29 @@ test(
   }),
 );
 
+test(
+  'allows passing functions to runJS',
+  withBrowser(async ({ utils }) => {
+    const mockFuncA = jest.fn(() => 5);
+    const mockFuncB = jest.fn();
+
+    await utils.runJS(
+      `
+      export default async (mockFuncA, mockFuncB) => {
+        const val = await mockFuncA('hello world')
+        if (val !== 5) throw new Error('Did not get return value');
+        await mockFuncB()
+      }
+    `,
+      [mockFuncA, mockFuncB],
+    );
+
+    expect(mockFuncA).toHaveBeenCalledTimes(1);
+    expect(mockFuncA).toHaveBeenCalledWith('hello world');
+    expect(mockFuncB).toHaveBeenCalledTimes(1);
+  }),
+);
+
 test(
   'allows passing ElementHandles and serializable values into browser',
   withBrowser(async ({ utils, screen }) => {

From 6562b6767874edeb78aa40cb19a1193103c35fdf Mon Sep 17 00:00:00 2001
From: Gerardo Rodriguez <mr.gerardo.rodriguez@gmail.com>
Date: Tue, 31 Aug 2021 16:18:28 -0700
Subject: [PATCH 2/3] New feature logic

---
 src/index.ts | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index cf74c501..b9326c29 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -297,11 +297,30 @@ const createTab = async ({
 
   await page.goto(`http://localhost:${port}`);
 
+  // eslint-disable-next-line @cloudfour/typescript-eslint/ban-types
+  const functionArgs: Function[] = [];
+
   const runJS: PleasantestUtils['runJS'] = (code, args) =>
     asyncHookTracker.addHook(async () => {
+      await page
+        .exposeFunction('pleasantest_callFunction', (id, args) =>
+          functionArgs[id](...args),
+        )
+        .catch((error) => {
+          if (!error.message.includes('already exists')) throw error;
+        });
       // For some reason encodeURIComponent doesn't encode '
       const encodedCode = encodeURIComponent(code).replace(/'/g, '%27');
       const buildStatus = createBuildStatusTracker();
+
+      const argsWithFuncsAsObjs = args?.map((arg) => {
+        if (typeof arg === 'function') {
+          const id = functionArgs.push(arg) - 1;
+          return { isFunction: true, id };
+        }
+        return arg;
+      });
+
       // This uses the testPath as the url so that if there are relative imports
       // in the inline code, the relative imports are resolved relative to the test file
       const url = `http://localhost:${port}/${testPath}?inline-code=${encodedCode}&build-id=${buildStatus.buildId}`;
@@ -310,14 +329,24 @@ const createTab = async ({
           '...args',
           `return import(${JSON.stringify(url)})
             .then(async m => {
-              if (m.default) await m.default(...args)
+              const argsWithFuncs = args.map(arg => {
+                if (typeof arg === 'object' && arg && arg.isFunction) {
+                  return async (...args) => {
+                    return await window.pleasantest_callFunction(arg.id, args);
+                  }
+                }
+                return arg
+              })
+              if (m.default) await m.default(...argsWithFuncs)
             })
             .catch(e =>
              e instanceof Error
                ? { message: e.message, stack: e.stack }
                : e)`,
         ) as () => any,
-        ...(Array.isArray(args) ? (args as any) : []),
+        ...(Array.isArray(argsWithFuncsAsObjs)
+          ? (argsWithFuncsAsObjs as any)
+          : []),
       );
 
       const errorsFromBuild = buildStatus.complete();

From bf53e310b0480eeaf19451c5cde02d19dfb4edd1 Mon Sep 17 00:00:00 2001
From: Gerardo Rodriguez <mr.gerardo.rodriguez@gmail.com>
Date: Tue, 31 Aug 2021 16:25:25 -0700
Subject: [PATCH 3/3] Add changeset

---
 .changeset/rotten-news-protect.md | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 .changeset/rotten-news-protect.md

diff --git a/.changeset/rotten-news-protect.md b/.changeset/rotten-news-protect.md
new file mode 100644
index 00000000..5f231ece
--- /dev/null
+++ b/.changeset/rotten-news-protect.md
@@ -0,0 +1,5 @@
+---
+'pleasantest': minor
+---
+
+Allow functions to be passed to runJS