Workers and Messaging

/**
 * Worker job configuration
 */
interface Job {
  /** Test capabilities */
  caps: WebdriverIO.Capabilities;
  /** Test specification files */
  specs: string[];
  /** Whether job has tests to run */
  hasTests: boolean;
  /** Base URL for tests */
  baseUrl?: string;
  /** Test runner configuration */
  config?: TestrunnerOptions & { sessionId?: string };
  /** Resolved capabilities */
  capabilities?: WebdriverIO.Capabilities;
}

/**
 * Worker message arguments (subset of Job)
 */
type WorkerMessageArgs = Omit<Job, 'caps' | 'specs' | 'hasTests'>;

/**
 * Worker execution payload
 */
interface WorkerRunPayload {
  /** Capability ID */
  cid: string;
  /** Configuration file path */
  configFile: string;
  /** Test capabilities */
  caps: WebdriverIO.Capabilities;
  /** Test specification files */
  specs: string[];
  /** Node.js execution arguments */
  execArgv: string[];
  /** Number of retry attempts */
  retries: number;
}

/**
 * Worker command structure
 */
interface WorkerCommand extends Omit<WorkerRunPayload, 'execArgv'> {
  /** Command to execute */
  command: string;
  /** Command arguments */
  args: any;
}

/**
 * Worker instance interface
 */
interface Worker extends Omit<TestrunnerOptions, 'capabilities' | 'specs' | 'rootDir'>, EventEmitter {
  /** Worker capabilities */
  capabilities: WebdriverIO.Capabilities;
  /** Worker configuration */
  config: TestrunnerOptions;
  /** Resolved capabilities */
  caps: WebdriverIO.Capabilities;
  /** Capability ID */
  cid: string;
  /** Whether worker is currently busy */
  isBusy?: boolean;
  /** Send message to worker */
  postMessage: (command: string, args: WorkerMessageArgs) => void;
  /** Test specification files */
  specs: string[];
  /** WebDriver session ID */
  sessionId?: string;
  /** Aggregated log messages */
  logsAggregator: string[];
}

/**
 * Worker pool type
 */
type WorkerPool = Record<string, Worker>;

/**
 * Message types for worker communication
 */
enum MESSAGE_TYPES {
  // Browser runner messages
  consoleMessage = 0,
  commandRequestMessage,
  commandResponseMessage,
  hookTriggerMessage,
  hookResultMessage,
  expectRequestMessage,
  expectResponseMessage,
  expectMatchersRequest,
  expectMatchersResponse,
  coverageMap,
  customCommand,
  initiateBrowserStateRequest,
  initiateBrowserStateResponse,
  browserTestResult
}

/**
 * Socket message value mapping
 */
type SocketMessageValue = {
  [MESSAGE_TYPES.consoleMessage]: ConsoleEvent;
  [MESSAGE_TYPES.commandRequestMessage]: CommandRequestEvent;
  [MESSAGE_TYPES.commandResponseMessage]: CommandResponseEvent;
  [MESSAGE_TYPES.hookTriggerMessage]: HookTriggerEvent;
  [MESSAGE_TYPES.hookResultMessage]: HookResultEvent;
  [MESSAGE_TYPES.expectRequestMessage]: ExpectRequestEvent;
  [MESSAGE_TYPES.expectResponseMessage]: ExpectResponseEvent;
  [MESSAGE_TYPES.expectMatchersRequest]: never;
  [MESSAGE_TYPES.expectMatchersResponse]: ExpectMatchersResponse;
  [MESSAGE_TYPES.coverageMap]: any;
  [MESSAGE_TYPES.customCommand]: CustomCommandEvent;
  [MESSAGE_TYPES.initiateBrowserStateRequest]: BrowserStateRequest;
  [MESSAGE_TYPES.initiateBrowserStateResponse]: BrowserState;
  [MESSAGE_TYPES.browserTestResult]: BrowserTestResults;
};

/**
 * Generic socket message payload
 */
type SocketMessagePayload<T extends MESSAGE_TYPES> = T extends any
  ? { type: T; value: SocketMessageValue[T] }
  : never;

/**
 * Union of all socket message types
 */
type SocketMessage = SocketMessagePayload<MESSAGE_TYPES>;

/**
 * Console logging event
 */
interface ConsoleEvent {
  /** Event name */
  name: 'consoleEvent';
  /** Log type */
  type: 'log' | 'info' | 'warn' | 'debug' | 'error';
  /** Log arguments */
  args: unknown[];
  /** Capability ID */
  cid: string;
}

/**
 * Browser test execution results
 */
interface BrowserTestResults {
  /** Number of test failures */
  failures: number;
  /** Test events */
  events: any[];
}

/**
 * Custom command event
 */
interface CustomCommandEvent {
  /** Command name */
  commandName: string;
  /** Capability ID */
  cid: string;
}

/**
 * Browser state request
 */
interface BrowserStateRequest {
  /** Capability ID */
  cid: string;
}

/**
 * Browser state response
 */
interface BrowserState {
  /** Available custom commands */
  customCommands: string[];
}

/**
 * Available matchers response
 */
interface ExpectMatchersResponse {
  /** List of available matchers */
  matchers: string[];
}

/**
 * Base interface for messages with pending promise ID
 */
interface MessageWithPendingPromiseId {
  /** Unique message identifier */
  id: number;
}

/**
 * Hook trigger event
 */
interface HookTriggerEvent extends MessageWithPendingPromiseId {
  /** Capability ID */
  cid: string;
  /** Hook name */
  name: string;
  /** Hook arguments */
  args: unknown[];
}

/**
 * Hook execution result
 */
interface HookResultEvent extends MessageWithPendingPromiseId {
  /** Error if hook failed */
  error?: Error;
}

/**
 * WebDriver command request
 */
interface CommandRequestEvent extends MessageWithPendingPromiseId {
  /** Capability ID */
  cid: string;
  /** WebDriver command name */
  commandName: string;
  /** Command arguments */
  args: unknown[];
  /** Command scope (browser, element, etc.) */
  scope?: string;
}

/**
 * WebDriver command response
 */
interface CommandResponseEvent extends MessageWithPendingPromiseId {
  /** Command result */
  result?: unknown;
  /** Error if command failed */
  error?: Error;
}

/**
 * Assertion/expectation request
 */
interface ExpectRequestEvent extends MessageWithPendingPromiseId {
  /** Capability ID */
  cid: string;
  /** Matcher name */
  matcherName: string;
  /** Matcher state (from expect library) */
  scope: any;
  /** Matcher arguments */
  args: unknown[];
  /** Element(s) being asserted */
  element?: any | any[];
  /** Additional context */
  context?: unknown;
  /** Error stack for inline snapshots */
  errorStack?: string;
}

/**
 * Assertion/expectation result
 */
interface ExpectResponseEvent extends MessageWithPendingPromiseId {
  /** Whether assertion passed */
  pass: boolean;
  /** Assertion message */
  message: string;
}

/**
 * Worker request message
 */
interface WorkerRequest {
  /** Message type */
  command: 'workerRequest';
  /** Request arguments */
  args: {
    /** Request ID */
    id: number;
    /** Socket message */
    message: SocketMessage;
  };
}

/**
 * Worker event message
 */
interface WorkerEvent {
  /** Message type */
  name: 'workerEvent';
  /** Event origin */
  origin: string;
  /** Socket message arguments */
  args: SocketMessage;
}

/**
 * Generic worker message
 */
interface WorkerMessage {
  /** Message name */
  name: string;
  /** Message content */
  content: {
    /** Session ID */
    sessionId?: string;
    /** Multiremote flag */
    isMultiremote?: boolean;
    /** Capabilities */
    capabilities: WebdriverIO.Capabilities;
  };
  /** Message origin */
  origin: string;
  /** Message parameters */
  params: Record<string, string>;
}

import type { Workers } from "@wdio/types";
import { EventEmitter } from "node:events";

// Custom worker implementation
class CustomWorker extends EventEmitter implements Workers.Worker {
  public capabilities: WebdriverIO.Capabilities;
  public config: any;
  public caps: WebdriverIO.Capabilities;
  public cid: string;
  public specs: string[];
  public sessionId?: string;
  public logsAggregator: string[] = [];
  public isBusy: boolean = false;

  constructor(
    cid: string,
    capabilities: WebdriverIO.Capabilities,
    specs: string[],
    config: any
  ) {
    super();
    this.cid = cid;
    this.capabilities = capabilities;
    this.caps = capabilities;
    this.specs = specs;
    this.config = config;
  }

  postMessage(command: string, args: Workers.WorkerMessageArgs) {
    console.log(`Worker ${this.cid} received command: ${command}`);
    
    // Handle different message types
    switch (command) {
      case 'run':
        this.runTests(args);
        break;
      case 'stop':
        this.stopTests();
        break;
      default:
        console.warn(`Unknown command: ${command}`);
    }
  }

  private runTests(args: Workers.WorkerMessageArgs) {
    this.isBusy = true;
    console.log(`Running tests: ${this.specs.join(', ')}`);
    
    // Simulate test execution
    setTimeout(() => {
      this.emit('test:complete', {
        cid: this.cid,
        specs: this.specs,
        passed: Math.random() > 0.5
      });
      this.isBusy = false;
    }, 1000);
  }

  private stopTests() {
    this.isBusy = false;
    console.log(`Stopping tests for worker ${this.cid}`);
  }
}

// Message handler for different message types
class MessageHandler {
  handleMessage(message: Workers.SocketMessage) {
    switch (message.type) {
      case Workers.MESSAGE_TYPES.consoleMessage:
        this.handleConsoleMessage(message.value);
        break;
      
      case Workers.MESSAGE_TYPES.commandRequestMessage:
        this.handleCommandRequest(message.value);
        break;
      
      case Workers.MESSAGE_TYPES.commandResponseMessage:
        this.handleCommandResponse(message.value);
        break;
      
      case Workers.MESSAGE_TYPES.expectRequestMessage:
        this.handleExpectRequest(message.value);
        break;
      
      case Workers.MESSAGE_TYPES.expectResponseMessage:
        this.handleExpectResponse(message.value);
        break;
      
      default:
        console.log(`Unhandled message type: ${message.type}`);
    }
  }

  private handleConsoleMessage(event: Workers.ConsoleEvent) {
    console.log(`[${event.cid}] ${event.type.toUpperCase()}:`, ...event.args);
  }

  private handleCommandRequest(event: Workers.CommandRequestEvent) {
    console.log(`Command request: ${event.commandName}`, event.args);
  }

  private handleCommandResponse(event: Workers.CommandResponseEvent) {
    if (event.error) {
      console.error(`Command failed:`, event.error);
    } else {
      console.log(`Command result:`, event.result);
    }
  }

  private handleExpectRequest(event: Workers.ExpectRequestEvent) {
    console.log(`Assertion: ${event.matcherName}`, event.args);
  }

  private handleExpectResponse(event: Workers.ExpectResponseEvent) {
    const status = event.pass ? 'PASS' : 'FAIL';
    console.log(`${status}: ${event.message}`);
  }
}

// Worker pool management
class WorkerPool {
  private workers: Workers.WorkerPool = {};

  addWorker(worker: Workers.Worker) {
    this.workers[worker.cid] = worker;
    
    worker.on('test:complete', (result) => {
      console.log(`Worker ${worker.cid} completed tests:`, result);
    });
  }

  removeWorker(cid: string) {
    const worker = this.workers[cid];
    if (worker) {
      worker.removeAllListeners();
      delete this.workers[cid];
    }
  }

  getWorker(cid: string): Workers.Worker | undefined {
    return this.workers[cid];
  }

  getAllWorkers(): Workers.Worker[] {
    return Object.values(this.workers);
  }

  getAvailableWorkers(): Workers.Worker[] {
    return Object.values(this.workers).filter(worker => !worker.isBusy);
  }
}