feat: initial support for test scenarios

Author: urish

package-lock.json

@@ -22,7 +22,8 @@
   "dependencies": {
     "@iarna/toml": "2.2.5",
     "arg": "^5.0.2",
-    "ws": "^8.13.0"
+    "ws": "^8.13.0",
+    "yaml": "^2.3.1"
   },
   "devDependencies": {
     "@types/ws": "^8.5.4",

package.json

@@ -9,11 +9,13 @@ import type {
   APISimStartParams,
 } from './APITypes';
 
-const DEFAULT_SERVER = 'wss://wokwi.com/api/ws/beta';
+const DEFAULT_SERVER = process.env.WOKWI_CLI_SERVER ?? 'wss://wokwi.com/api/ws/beta';
 
 export class APIClient {
   private readonly socket: WebSocket;
   private lastId = 0;
+  private _running = false;
+  private _lastNanos = 0;
   private readonly pendingCommands = new Map<
     string,
     [(result: any) => void, (error: Error) => void]
@@ -52,6 +54,7 @@ export class APIClient {
   }
 
   async simStart(params: APISimStartParams) {
+    this._running = false;
     return await this.sendCommand('sim:start', params);
   }
 
@@ -60,6 +63,7 @@ export class APIClient {
   }
 
   async simResume(pauseAfter?: number) {
+    this._running = true;
     return await this.sendCommand('sim:resume', { pauseAfter });
   }
 
@@ -83,6 +87,10 @@ export class APIClient {
     return await this.sendCommand<{ png: string }>('framebuffer:read', { id: partId });
   }
 
+  async controlSet(partId: string, control: string, value: number) {
+    return await this.sendCommand('control:set', { part: partId, control, value });
+  }
+
   async sendCommand<T = unknown>(command: string, params?: any) {
     return await new Promise<T>((resolve, reject) => {
       const id = this.lastId++;
@@ -92,6 +100,14 @@ export class APIClient {
     });
   }
 
+  get running() {
+    return this._running;
+  }
+
+  get lastNanos() {
+    return this._lastNanos;
+  }
+
   processMessage(message: APIError | APIHello | APIEvent | APIResponse) {
     switch (message.type) {
       case 'error':
@@ -116,6 +132,10 @@ export class APIClient {
   }
 
   processEvent(message: APIEvent) {
+    if (message.event === 'sim:pause') {
+      this._running = false;
+    }
+    this._lastNanos = message.nanos;
     this.onEvent?.(message);
   }
 

src/APIClient.ts

@@ -0,0 +1,140 @@
+import chalk from 'chalk';
+import type { APIClient } from './APIClient';
+import type { EventManager } from './EventManager';
+import type { ExpectEngine } from './ExpectEngine';
+import { parseTime } from './utils/parseTime';
+
+const validStepKeys = ['name'];
+
+export interface ISetControlParams {
+  'part-id': string;
+  control: string;
+  value: number;
+}
+
+export interface IStepDefinition {
+  name?: string;
+  'wait-serial': string;
+  delay: string;
+  'set-control': ISetControlParams;
+}
+
+export interface IScenarioDefinition {
+  name: string;
+  version: number;
+  author?: string;
+
+  steps: IStepDefinition[];
+}
+
+export class TestScenario {
+  private stepIndex = 0;
+  private client?: APIClient;
+
+  constructor(
+    readonly scenario: IScenarioDefinition,
+    readonly eventManager: EventManager,
+    readonly expectEngine: ExpectEngine
+  ) {}
+
+  validate() {
+    const { scenario } = this;
+    if (scenario.name == null) {
+      throw new Error(`Scenario name is missing`);
+    }
+
+    if (typeof scenario.name !== 'string') {
+      throw new Error(`Scenario name must be a string`);
+    }
+
+    if (scenario.version !== 1) {
+      throw new Error(`Unsupported scenario version: ${scenario.version}`);
+    }
+
+    if (!Array.isArray(scenario.steps)) {
+      throw new Error(`Scenario steps must be an array`);
+    }
+
+    for (const step of scenario.steps) {
+      if (typeof step !== 'object') {
+        throw new Error(`Scenario step must be an object`);
+      }
+
+      for (const key of Object.keys(step)) {
+        if (!validStepKeys.includes(key) && !Object.keys(this.handlers).includes(key)) {
+          throw new Error(`Invalid scenario step key: ${key}`);
+        }
+      }
+    }
+  }
+
+  async start(client: APIClient) {
+    this.stepIndex = 0;
+    this.client = client;
+    await this.nextStep();
+  }
+
+  async nextStep() {
+    if (this.client?.running) {
+      void this.client.simPause();
+    }
+
+    const step = this.scenario.steps[this.stepIndex];
+    if (step == null) {
+      this.log(chalk`{green Scenario completed successfully}`);
+      process.exit(0);
+    }
+    if (step.name) {
+      this.log(chalk`{gray Executing step:} {yellow ${step.name}`);
+    }
+    for (const key of Object.keys(this.handlers) as Array<keyof typeof this.handlers>) {
+      if (key in step) {
+        const value = step[key];
+        console.log('running handler for ' + key);
+        void this.handlers[key](value as any, step);
+        this.stepIndex++;
+        return;
+      }
+    }
+    console.error('Unknown key in step: ', step);
+    process.exit(1);
+  }
+
+  log(message: string) {
+    console.log(chalk`{cyan [${this.scenario.name}]}`, message);
+  }
+
+  async resume() {
+    await this.client?.simResume(
+      this.eventManager.timeToNextEvent >= 0 ? this.eventManager.timeToNextEvent : undefined
+    );
+  }
+
+  handlers = {
+    'wait-serial': async (text: string) => {
+      this.expectEngine.expectTexts.push(text);
+      this.expectEngine.once('match', () => {
+        this.log(chalk`Expected text matched: {green "${text}"}`);
+        const textIndex = this.expectEngine.expectTexts.indexOf(text);
+        if (textIndex >= 0) {
+          this.expectEngine.expectTexts.splice(textIndex, 1);
+        }
+        void this.nextStep();
+      });
+      await this.resume();
+    },
+    delay: async (value: string, step: IStepDefinition) => {
+      const nanos = parseTime(value);
+      const targetNanos = (this.client?.lastNanos ?? 0) + nanos;
+      this.log(chalk`delay {yellow "${value}"}`);
+      this.eventManager.at(targetNanos, () => {
+        void this.nextStep();
+      });
+      await this.resume();
+    },
+    'set-control': async (params: ISetControlParams) => {
+      await this.client?.controlSet(params['part-id'], params.control, params.value);
+      await this.nextStep();
+    },
+  };
+}

src/TestScenario.ts

@@ -2,14 +2,16 @@ import arg from 'arg';
 import chalk from 'chalk';
 import { existsSync, readFileSync, writeFileSync } from 'fs';
 import path, { join } from 'path';
+import YAML from 'yaml';
 import { APIClient } from './APIClient';
 import type { APIEvent, ChipsLogPayload, SerialMonitorDataPayload } from './APITypes';
 import { EventManager } from './EventManager';
 import { ExpectEngine } from './ExpectEngine';
+import { TestScenario } from './TestScenario';
 import { parseConfig } from './config';
 import { cliHelp } from './help';
-import { readVersion } from './readVersion';
 import { loadChips } from './loadChips';
+import { readVersion } from './readVersion';
 
 const millis = 1_000_000;
 
@@ -21,6 +23,7 @@ async function main() {
       '--version': Boolean,
       '--expect-text': String,
       '--fail-text': String,
+      '--scenario': String,
       '--screenshot-part': String,
       '--screenshot-file': String,
       '--screenshot-time': Number,
@@ -35,6 +38,7 @@ async function main() {
   const quiet = args['--quiet'];
   const expectText = args['--expect-text'];
   const failText = args['--fail-text'];
+  const scenarioFile = args['--scenario'];
   const timeout = args['--timeout'] ?? 0;
   const screenshotPart = args['--screenshot-part'];
   const screenshotTime = args['--screenshot-time'];
@@ -93,11 +97,31 @@ async function main() {
 
   const chips = loadChips(config.chip ?? [], rootDir);
 
+  if (scenarioFile && !existsSync(scenarioFile)) {
+    console.error(`Error: scenario file not found: ${path.resolve(scenarioFile)}`);
+    process.exit(1);
+  }
+
+  const eventManager = new EventManager();
   const expectEngine = new ExpectEngine();
 
+  let scenario;
+  if (scenarioFile) {
+    scenario = new TestScenario(
+      YAML.parse(readFileSync(scenarioFile, 'utf-8')),
+      eventManager,
+      expectEngine
+    );
+    scenario.validate();
+  }
+
   if (expectText) {
     expectEngine.expectTexts.push(expectText);
     expectEngine.on('match', (text) => {
+      if (text !== expectText) {
+        return;
+      }
+
       if (!quiet) {
         console.log(chalk`\n\nExpected text found: {green "${expectText}"}`);
         console.log('TEST PASSED.');
@@ -109,14 +133,17 @@ async function main() {
   if (failText) {
     expectEngine.failTexts.push(failText);
     expectEngine.on('fail', (text) => {
+      if (text !== failText) {
+        return;
+      }
+
       console.error(chalk`\n\n{red Error:} Unexpected text found: {yellow "${text}"}`);
       console.error('TEST FAILED.');
       process.exit(1);
     });
   }
 
   const client = new APIClient(token);
-  const eventManager = new EventManager();
   client.onConnected = (hello) => {
     if (!quiet) {
       console.log(`Connected to Wokwi Simulation API ${hello.appVersion}`);
@@ -138,6 +165,8 @@ async function main() {
     console.log('Starting simulation...');
   }
 
+  await scenario?.start(client);
+
   if (timeoutNanos) {
     eventManager.at(timeoutNanos, () => {
       // We are using setImmediate to make sure other events (e.g. screen shot) are processed first