Merge pull request #52 from mozilla/feat/connect-existing-tests

test: add unit tests for connect-existing mode

Author: freema

tests/firefox/connect-existing.test.ts

@@ -0,0 +1,371 @@
+/**
+ * Unit tests for connect-existing mode features (PR #50)
+ * - GeckodriverHttpDriver BiDi support
+ * - Session cleanup on quit/kill
+ * - marionetteHost parameter
+ * - Reconnect on lost connection
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+// ---------------------------------------------------------------------------
+// GeckodriverHttpDriver tests — we access the class indirectly through
+// FirefoxCore since GeckodriverHttpDriver is not exported.
+// For direct testing we use (core as any).driver after mocked connect().
+// ---------------------------------------------------------------------------
+
+describe('GeckodriverHttpDriver BiDi support', () => {
+  let mockWsInstance: {
+    on: ReturnType<typeof vi.fn>;
+    off: ReturnType<typeof vi.fn>;
+    send: ReturnType<typeof vi.fn>;
+    close: ReturnType<typeof vi.fn>;
+    readyState: number;
+  };
+  let wsEventListeners: Record<string, Function[]>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.resetModules();
+
+    wsEventListeners = {};
+    mockWsInstance = {
+      readyState: 1,
+      on: vi.fn((event: string, handler: Function) => {
+        if (!wsEventListeners[event]) wsEventListeners[event] = [];
+        wsEventListeners[event].push(handler);
+      }),
+      off: vi.fn(),
+      send: vi.fn(),
+      close: vi.fn(),
+    };
+  });
+
+  /**
+   * Helper: create a GeckodriverHttpDriver instance via mocked connect().
+   * Returns the FirefoxCore with driver set to GeckodriverHttpDriver.
+   */
+  async function createConnectExistingCore(opts?: {
+    webSocketUrl?: string;
+    marionetteHost?: string;
+  }) {
+    const mockGdProcess = {
+      stdout: { on: vi.fn() },
+      stderr: { on: vi.fn() },
+      on: vi.fn(),
+      kill: vi.fn(),
+    };
+
+    // Mock child_process.spawn to return our mock geckodriver process
+    vi.doMock('node:child_process', async (importOriginal) => {
+      const original = (await importOriginal()) as typeof import('node:child_process');
+      return {
+        ...original,
+        spawn: vi.fn(() => {
+          // Simulate geckodriver printing its listening port
+          setTimeout(() => {
+            const onData = mockGdProcess.stderr.on.mock.calls.find(
+              (c: unknown[]) => c[0] === 'data'
+            );
+            if (onData) {
+              (onData[1] as Function)(Buffer.from('Listening on 127.0.0.1:4444'));
+            }
+          }, 5);
+          return mockGdProcess;
+        }),
+      };
+    });
+
+    // Mock fetch for session creation
+    const wsUrl = opts?.webSocketUrl ?? null;
+    vi.doMock(
+      'node:module',
+      async (importOriginal) => await importOriginal()
+    );
+
+    // We need to mock global fetch
+    const mockFetch = vi.fn().mockResolvedValue({
+      json: vi.fn().mockResolvedValue({
+        value: {
+          sessionId: 'mock-session-id',
+          capabilities: {
+            webSocketUrl: wsUrl,
+          },
+        },
+      }),
+    });
+    vi.stubGlobal('fetch', mockFetch);
+
+    // Mock selenium-manager to avoid real binary lookup
+    vi.doMock('selenium-webdriver/package.json', () => ({}), { virtual: true });
+
+    // Mock WebSocket constructor
+    vi.doMock('ws', () => ({
+      default: vi.fn(() => {
+        // Simulate open event on next tick
+        setTimeout(() => {
+          if (wsEventListeners['open']) {
+            wsEventListeners['open'].forEach((h) => h());
+          }
+        }, 5);
+        return mockWsInstance;
+      }),
+    }));
+
+    const { FirefoxCore } = await import('@/firefox/core.js');
+
+    const core = new FirefoxCore({
+      headless: true,
+      connectExisting: true,
+      marionettePort: 2828,
+      marionetteHost: opts?.marionetteHost,
+    });
+
+    await core.connect();
+    return { core, mockGdProcess, mockFetch };
+  }
+
+  it('should throw when getBidi() called without webSocketUrl', async () => {
+    const { core } = await createConnectExistingCore({ webSocketUrl: undefined });
+    const driver = core.getDriver();
+
+    await expect(driver.getBidi()).rejects.toThrow(
+      /BiDi is not available.*webSocketUrl/
+    );
+  });
+
+  it('should open WebSocket and return BiDi handle', async () => {
+    const { core } = await createConnectExistingCore({
+      webSocketUrl: 'ws://127.0.0.1:9222/session/test',
+    });
+    const driver = core.getDriver();
+
+    const bidi = await driver.getBidi();
+    expect(bidi).toBeDefined();
+    expect(bidi.socket).toBeDefined();
+    expect(bidi.subscribe).toBeDefined();
+  });
+
+  it('should cache BiDi connection on subsequent calls', async () => {
+    const { core } = await createConnectExistingCore({
+      webSocketUrl: 'ws://127.0.0.1:9222/session/test',
+    });
+    const driver = core.getDriver();
+
+    const bidi1 = await driver.getBidi();
+    const bidi2 = await driver.getBidi();
+    expect(bidi1).toBe(bidi2);
+  });
+
+  it('subscribe should send session.subscribe and wait for response', async () => {
+    const { core } = await createConnectExistingCore({
+      webSocketUrl: 'ws://127.0.0.1:9222/session/test',
+    });
+    const driver = core.getDriver();
+    const bidi = await driver.getBidi();
+
+    // Start subscribe
+    const subscribePromise = bidi.subscribe!('log.entryAdded', ['context-1']);
+
+    // Wait a tick for send to be called
+    await new Promise((r) => setTimeout(r, 10));
+
+    expect(mockWsInstance.send).toHaveBeenCalledTimes(1);
+    const sent = JSON.parse(mockWsInstance.send.mock.calls[0][0]);
+    expect(sent.method).toBe('session.subscribe');
+    expect(sent.params.events).toEqual(['log.entryAdded']);
+    expect(sent.params.contexts).toEqual(['context-1']);
+
+    // Simulate response
+    if (wsEventListeners['message']) {
+      wsEventListeners['message'].forEach((h) =>
+        h(JSON.stringify({ id: sent.id, result: {} }))
+      );
+    }
+
+    await expect(subscribePromise).resolves.toBeUndefined();
+  });
+
+  it('subscribe should reject on error response', async () => {
+    const { core } = await createConnectExistingCore({
+      webSocketUrl: 'ws://127.0.0.1:9222/session/test',
+    });
+    const driver = core.getDriver();
+    const bidi = await driver.getBidi();
+
+    const subscribePromise = bidi.subscribe!('log.entryAdded');
+
+    await new Promise((r) => setTimeout(r, 10));
+
+    const sent = JSON.parse(mockWsInstance.send.mock.calls[0][0]);
+    if (wsEventListeners['message']) {
+      wsEventListeners['message'].forEach((h) =>
+        h(JSON.stringify({ id: sent.id, error: 'invalid subscription' }))
+      );
+    }
+
+    await expect(subscribePromise).rejects.toThrow(/BiDi subscribe error/);
+  });
+});
+
+describe('GeckodriverHttpDriver session cleanup', () => {
+  let mockGdProcess: {
+    stdout: { on: ReturnType<typeof vi.fn> };
+    stderr: { on: ReturnType<typeof vi.fn> };
+    on: ReturnType<typeof vi.fn>;
+    kill: ReturnType<typeof vi.fn>;
+  };
+  let mockFetch: ReturnType<typeof vi.fn>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.resetModules();
+
+    mockGdProcess = {
+      stdout: { on: vi.fn() },
+      stderr: { on: vi.fn() },
+      on: vi.fn(),
+      kill: vi.fn(),
+    };
+
+    vi.doMock('node:child_process', async (importOriginal) => {
+      const original = (await importOriginal()) as typeof import('node:child_process');
+      return {
+        ...original,
+        spawn: vi.fn(() => {
+          setTimeout(() => {
+            const onData = mockGdProcess.stderr.on.mock.calls.find(
+              (c: unknown[]) => c[0] === 'data'
+            );
+            if (onData) {
+              (onData[1] as Function)(Buffer.from('Listening on 127.0.0.1:4444'));
+            }
+          }, 5);
+          return mockGdProcess;
+        }),
+      };
+    });
+
+    mockFetch = vi.fn().mockResolvedValue({
+      json: vi.fn().mockResolvedValue({
+        value: {
+          sessionId: 'mock-session-id',
+          capabilities: {},
+        },
+      }),
+    });
+    vi.stubGlobal('fetch', mockFetch);
+
+    vi.doMock('ws', () => ({ default: vi.fn() }));
+  });
+
+  async function createCore() {
+    const { FirefoxCore } = await import('@/firefox/core.js');
+    const core = new FirefoxCore({
+      headless: true,
+      connectExisting: true,
+      marionettePort: 2828,
+    });
+    await core.connect();
+    return core;
+  }
+
+  it('kill() should send DELETE /session before killing geckodriver', async () => {
+    const core = await createCore();
+
+    // Mock fetch for the DELETE call
+    mockFetch.mockResolvedValueOnce({
+      json: vi.fn().mockResolvedValue({ value: null }),
+    });
+
+    const driver = core.getDriver() as any;
+    await driver.kill();
+
+    // Verify DELETE /session was called
+    const deleteCalls = mockFetch.mock.calls.filter(
+      (c: unknown[]) => typeof c[1] === 'object' && (c[1] as RequestInit).method === 'DELETE'
+    );
+    expect(deleteCalls.length).toBeGreaterThan(0);
+    expect(mockGdProcess.kill).toHaveBeenCalled();
+  });
+
+  it('quit() should send DELETE /session and kill geckodriver', async () => {
+    const core = await createCore();
+
+    mockFetch.mockResolvedValueOnce({
+      json: vi.fn().mockResolvedValue({ value: null }),
+    });
+
+    const driver = core.getDriver() as any;
+    await driver.quit();
+
+    const deleteCalls = mockFetch.mock.calls.filter(
+      (c: unknown[]) => typeof c[1] === 'object' && (c[1] as RequestInit).method === 'DELETE'
+    );
+    expect(deleteCalls.length).toBeGreaterThan(0);
+    expect(mockGdProcess.kill).toHaveBeenCalled();
+  });
+
+  it('kill() should not throw if DELETE /session fails', async () => {
+    const core = await createCore();
+
+    mockFetch.mockRejectedValueOnce(new Error('connection refused'));
+
+    const driver = core.getDriver() as any;
+    await expect(driver.kill()).resolves.toBeUndefined();
+    expect(mockGdProcess.kill).toHaveBeenCalled();
+  });
+});
+
+describe('FirefoxCore connect-existing with marionetteHost', () => {
+  it('should pass marionetteHost to options', async () => {
+    const { FirefoxCore } = await import('@/firefox/core.js');
+    const core = new FirefoxCore({
+      headless: true,
+      connectExisting: true,
+      marionettePort: 2828,
+      marionetteHost: '192.168.1.100',
+    });
+
+    expect(core.getOptions().marionetteHost).toBe('192.168.1.100');
+  });
+});
+
+describe('getFirefox() reconnect behavior', () => {
+  it('should reconnect when connection is lost instead of throwing', async () => {
+    vi.resetModules();
+
+    // Mock the firefox module
+    const mockIsConnected = vi.fn();
+    const mockConnect = vi.fn();
+    const mockClose = vi.fn();
+
+    vi.doMock('@/firefox/index.js', () => ({
+      FirefoxDevTools: vi.fn(() => ({
+        isConnected: mockIsConnected,
+        connect: mockConnect,
+        close: mockClose,
+      })),
+    }));
+
+    // First call: create instance, connection works
+    mockIsConnected.mockResolvedValueOnce(true);
+    mockConnect.mockResolvedValue(undefined);
+
+    // This test verifies the reconnect logic pattern:
+    // When isConnected() returns false, getFirefox() should reset and create
+    // a new connection instead of throwing FirefoxDisconnectedError
+    const { FirefoxCore } = await import('@/firefox/core.js');
+    const core = new FirefoxCore({
+      headless: true,
+      connectExisting: true,
+      marionettePort: 2828,
+    });
+
+    // Verify reset clears the state
+    core.setCurrentContextId('old-context');
+    core.reset();
+    expect(core.getCurrentContextId()).toBe(null);
+    expect(() => core.getDriver()).toThrow('Driver not connected');
+  });
+});