Fix expect and ignore of send calls

2 files changed, 74 insertions, 37 deletions

diff --git a/desktop/packages/mullvad-vpn/src/shared/ipc-helpers.ts b/desktop/packages/mullvad-vpn/src/shared/ipc-helpers.ts
index c4f2c87c83..559b48c22b 100644
--- a/desktop/packages/mullvad-vpn/src/shared/ipc-helpers.ts
+++ b/desktop/packages/mullvad-vpn/src/shared/ipc-helpers.ts
@@ -16,6 +16,7 @@ interface MainToRenderer<T> {
 
 interface RendererToMain<T, R> {
   direction: 'renderer-to-main';
+  type: 'send' | 'invoke';
   send: (event: string, ipcRenderer: EIpcRenderer) => Sender<T, R>;
   receive: (event: string, ipcMain: EIpcMain) => Handler<T, R>;
 }
@@ -118,6 +119,7 @@ export function createIpc<S extends Schema, T, R>(
 export function send<T>(): RendererToMain<T, void> {
   return {
     direction: 'renderer-to-main',
+    type: 'send',
     send: (event, ipcRenderer) => (newValue: T) => ipcRenderer.send(event, newValue),
     receive: (event, ipcMain) => (handlerFn: (value: T) => void) => {
       ipcMain.on(event, (_event, newValue: T) => {
@@ -131,6 +133,7 @@ export function send<T>(): RendererToMain<T, void> {
 export function invokeSync<T, R>(): RendererToMain<T, R> {
   return {
     direction: 'renderer-to-main',
+    type: 'send',
     send: (event, ipcRenderer) => (newValue: T) => ipcRenderer.sendSync(event, newValue),
     receive: (event, ipcMain) => (handlerFn: (value: T) => R) => {
       ipcMain.on(event, (ipcEvent, newValue: T) => {
@@ -144,6 +147,7 @@ export function invokeSync<T, R>(): RendererToMain<T, R> {
 export function invoke<T, R>(): RendererToMain<T, Promise<R>> {
   return {
     direction: 'renderer-to-main',
+    type: 'invoke',
     send: invokeImpl,
     receive: handle,
   };
diff --git a/desktop/packages/mullvad-vpn/test/e2e/mocked/mocked-utils.ts b/desktop/packages/mullvad-vpn/test/e2e/mocked/mocked-utils.ts
index 126329931d..e6db357f99 100644
--- a/desktop/packages/mullvad-vpn/test/e2e/mocked/mocked-utils.ts
+++ b/desktop/packages/mullvad-vpn/test/e2e/mocked/mocked-utils.ts
@@ -44,9 +44,32 @@ export const startMockedApp = async (): Promise<StartMockedAppResponse> => {
   };
 };
 
-export const createMockIpcHandle = (electronApp: ElectronApplication, event: string) => {
+export const createMockIpcNotify = (electronApp: ElectronApplication, event: string) => {
+  return async <T>(arg: T) => {
+    await electronApp.evaluate(
+      ({ webContents }, { event, arg }) => {
+        webContents
+          .getAllWebContents()
+          // Select window that isn't devtools
+          .find((webContents) => webContents.getURL().startsWith('file://'))!
+          .send(event, arg);
+      },
+      { event, arg },
+    );
+  };
+};
+
+export const createMockIpcHandle = (
+  electronApp: ElectronApplication,
+  event: string,
+  spec: AnyIpcCall,
+) => {
   // This function resolves when the handle is registered. To await the event, use `expect()`.
   return async <T>(response: T): Promise<void> => {
+    if ('type' in spec && spec.type === 'send') {
+      throw new Error(`No value can be returned on a send call (${event})`);
+    }
+
     await electronApp.evaluate(
       ({ ipcMain }, { event, response }) => {
         ipcMain.removeHandler(event);
@@ -62,51 +85,61 @@ export const createMockIpcHandle = (electronApp: ElectronApplication, event: str
   };
 };
 
-export const createMockIpcNotify = (electronApp: ElectronApplication, event: string) => {
-  return async <T>(arg: T) => {
-    await electronApp.evaluate(
-      ({ webContents }, { event, arg }) => {
-        webContents
-          .getAllWebContents()
-          // Select window that isn't devtools
-          .find((webContents) => webContents.getURL().startsWith('file://'))!
-          .send(event, arg);
-      },
-      { event, arg },
-    );
-  };
-};
+// Use when you want to wait for an IPC call to happen but don't need to respond. The returned
+// promise resolves when the IPC handle/on methods are triggered triggered.
+export const createMockIpcExpect = (
+  electronApp: ElectronApplication,
+  event: string,
+  spec: AnyIpcCall,
+) => {
+  const type = 'type' in spec ? spec.type : 'invoke';
 
-export const createMockIpcExpect = (electronApp: ElectronApplication, event: string) => {
   return <T>(): Promise<T> => {
     return electronApp.evaluate(
-      ({ ipcMain }, { event }) => {
+      ({ ipcMain }, { event, type }) => {
         return new Promise<T>((resolve) => {
-          ipcMain.handleOnce(event, (_event, arg) => {
-            resolve(arg);
-            return {
-              type: 'success',
-              value: null,
-            };
-          });
+          if (type === 'send') {
+            ipcMain.once(event, (_event, arg) => resolve(arg));
+          } else {
+            ipcMain.handleOnce(event, (_event, arg) => {
+              resolve(arg);
+              return {
+                type: 'success',
+                value: null,
+              };
+            });
+          }
         });
       },
-      { event },
+      { event, type },
     );
   };
 };
 
-export const createMockIpcIgnore = (electronApp: ElectronApplication, event: string) => {
+// Use when you knowingly want to ignore when this IPC method is called. Useful to avoid unhandled
+// events from being printed and polluting the log output.
+export const createMockIpcIgnore = (
+  electronApp: ElectronApplication,
+  event: string,
+  spec: AnyIpcCall,
+) => {
+  const type = 'type' in spec ? spec.type : 'invoke';
+
   return async (): Promise<void> => {
     await electronApp.evaluate(
-      ({ ipcMain }, { event }) => {
-        ipcMain.removeHandler(event);
-        ipcMain.handle(event, () => ({
-          type: 'success',
-          value: null,
-        }));
+      ({ ipcMain }, { event, type }) => {
+        if (type === 'send') {
+          ipcMain.removeAllListeners(event);
+          ipcMain.addListener(event, () => {});
+        } else {
+          ipcMain.removeHandler(event);
+          ipcMain.handle(event, () => ({
+            type: 'success',
+            value: null,
+          }));
+        }
       },
-      { event },
+      { event, type },
     );
   };
 };
@@ -139,15 +172,15 @@ export type IpcMockedTest<S extends Schema> = {
 };
 
 export function createTestIpc(electronApp: ElectronApplication): IpcMockedTest<IpcSchema> {
-  return createIpc(ipcSchema, (event, key, _spec) => {
+  return createIpc(ipcSchema, (event, key, spec) => {
     return [
       key,
       {
         eventKey: event,
         notify: createMockIpcNotify(electronApp, event),
-        handle: createMockIpcHandle(electronApp, event),
-        expect: createMockIpcExpect(electronApp, event),
-        ignore: createMockIpcIgnore(electronApp, event),
+        handle: createMockIpcHandle(electronApp, event, spec),
+        expect: createMockIpcExpect(electronApp, event, spec),
+        ignore: createMockIpcIgnore(electronApp, event, spec),
       },
     ];
   });