appium/packages/images-plugin/lib/plugin.ts

import type {MatchingOptions, OccurrenceOptions, SimilarityOptions} from '@appium/opencv';
import {util} from '@appium/support';
import type {ActionSequence, Element, ExternalDriver, MethodMap} from '@appium/types';
import {errors} from 'appium/driver.js';
import {BasePlugin} from 'appium/plugin.js';

import {compareImages} from './compare.js';
import {IMAGE_ELEMENT_PREFIX, IMAGE_STRATEGY} from './constants.js';
import {ImageElementFinder} from './finder.js';
import {ImageElement} from './image-element.js';

export class ImageElementPlugin extends BasePlugin {
  // this plugin supports a non-standard 'compare images' command
  static newMethodMap: MethodMap<ImageElementPlugin> = {
    '/session/:sessionId/appium/compare_images': {
      POST: {
        command: 'compareImages',
        payloadParams: {
          required: ['mode', 'firstImage', 'secondImage'],
          optional: ['options'],
        },
        neverProxy: true,
      },
    },
  } as const;
  readonly finder: ImageElementFinder;

  constructor(pluginName: string) {
    super(pluginName);
    this.finder = new ImageElementFinder();
  }

  async compareImages(
    next: () => Promise<any>,
    driver: ExternalDriver,
    mode: string,
    firstImage: string | Buffer,
    secondImage: string | Buffer,
    options?: MatchingOptions | SimilarityOptions | OccurrenceOptions,
  ): Promise<any> {
    return await compareImages(mode, firstImage, secondImage, options);
  }

  async findElement(next: () => Promise<any>, driver: ExternalDriver, ...args: any[]): Promise<any> {
    return await this._find(false, next, driver, ...args);
  }

  async findElements(next: () => Promise<any>, driver: ExternalDriver, ...args: any[]): Promise<any> {
    return await this._find(true, next, driver, ...args);
  }

  async handle(next: () => Promise<any>, driver: ExternalDriver, cmdName: string, ...args: any[]): Promise<any> {
    // if we have a command that involves an image element id, attempt to find the image element
    // and execute the command on it
    const imgElId = getImgElFromArgs(args);
    if (imgElId) {
      const imgEl = this.finder.getImageElement(imgElId);
      if (!imgEl) {
        throw new errors.NoSuchElementError();
      }
      return await ImageElement.execute(driver, imgEl, cmdName, ...args);
    }

    if (cmdName === 'deleteSession') {
      this.finder.clearImageElements();
    }

    // otherwise just do the normal thing
    return await next();
  }

  async performActions(
    next: () => Promise<any>,
    driver: ExternalDriver,
    actionSequences: ActionSequence[],
  ): Promise<any> {
    // Replace with coordinates when ActionSequence includes image elements.
    for (const actionSequence of actionSequences) {
      for (const action of actionSequence.actions) {
        // The actions that can have an Element as the origin are "pointerMove" and "scroll".
        if (!util.isPlainObject((action as any).origin)) {
          continue;
        }

        const actionWithEl = action as any;

        const elId = util.unwrapElement(actionWithEl.origin as Element);
        if (!elId?.startsWith(IMAGE_ELEMENT_PREFIX)) {
          continue;
        }

        const imgEl = this.finder.getImageElement(elId);
        if (!imgEl) {
          throw new errors.NoSuchElementError();
        }

        // Add the element's center to the offset. An unset offset must count as
        // zero, otherwise the resulting coordinate is NaN.
        actionWithEl.x = (actionWithEl.x ?? 0) + imgEl.center.x;
        actionWithEl.y = (actionWithEl.y ?? 0) + imgEl.center.y;
        // Set the origin to the viewport so that the external driver can process it using coordinates.
        delete actionWithEl.origin;
      }
    }

    return await next();
  }

  private async _find(
    multiple: boolean,
    next: () => Promise<any>,
    driver: ExternalDriver,
    ...args: any[]
  ): Promise<any> {
    const [strategy, selector] = args;

    // if we're not actually finding by image, just do the normal thing
    if (strategy !== IMAGE_STRATEGY) {
      return await next();
    }

    return await this.finder.findByImage(Buffer.from(selector, 'base64'), driver, {multiple});
  }
}

/**
 * Returns the first image-element id found in command args.
 * @param args Command arguments.
 */
export function getImgElFromArgs(args: any[]): string | undefined {
  return args.find((arg) => typeof arg === 'string' && arg.startsWith(IMAGE_ELEMENT_PREFIX));
}