fix(ct): correctly support custom element commands

src/browser/browser.js

@@ -6,6 +6,7 @@ const { SAVE_HISTORY_MODE } = require("../constants/config");
 const { X_REQUEST_ID_DELIMITER } = require("../constants/browser");
 const history = require("./history");
 const stacktrace = require("./stacktrace");
+const { getBrowserCommands, getElementCommands } = require("./history/commands");
 const addRunStepCommand = require("./commands/runStep").default;
 
 const CUSTOM_SESSION_OPTS = [
@@ -106,12 +107,17 @@ module.exports = class Browser {
     }
 
     _startCollectingCustomCommands() {
-        this._session.overwriteCommand("addCommand", (origCommand, name, ...rest) => {
-            if (!this._session[name]) {
-                this._customCommands.add(name);
+        const browserCommands = getBrowserCommands();
+        const elementCommands = getElementCommands();
+
+        this._session.overwriteCommand("addCommand", (origCommand, name, wrapper, elementScope, ...rest) => {
+            const isKnownCommand = elementScope ? elementCommands.includes(name) : browserCommands.includes(name);
+
+            if (!isKnownCommand) {
+                this._customCommands.add({ name, elementScope: Boolean(elementScope) });
             }
 
-            return origCommand(name, ...rest);
+            return origCommand(name, wrapper, elementScope, ...rest);
         });
     }
 
@@ -144,6 +150,7 @@ module.exports = class Browser {
     }
 
     get customCommands() {
-        return Array.from(this._customCommands);
+        const allCustomCommands = Array.from(this._customCommands);
+        return _.uniqWith(allCustomCommands, _.isEqual);
     }
 };

src/browser/history/commands.ts

@@ -8,6 +8,7 @@ const wdioBrowserCommands = [
     "$",
     "action",
     "actions",
+    "addCommand",
     "call",
     "custom$$",
     "custom$",
@@ -23,6 +24,7 @@ const wdioBrowserCommands = [
     "mockClearAll",
     "mockRestoreAll",
     "newWindow",
+    "overwriteCommand",
     "pause",
     "react$$",
     "react$",
@@ -54,8 +56,8 @@ const wdioElementCommands = [
     "dragAndDrop",
     "getAttribute",
     "getCSSProperty",
-    "getComputedRole",
     "getComputedLabel",
+    "getComputedRole",
     "getHTML",
     "getLocation",
     "getProperty",

src/browser/history/index.js

@@ -88,7 +88,7 @@ const overwriteBrowserCommands = (session, callstack) =>
     overwriteCommands({
         session,
         callstack,
-        commands: cmds.getBrowserCommands(),
+        commands: cmds.getBrowserCommands().filter(cmd => !shouldNotWrapCommand(cmd)),
         elementScope: false,
     });
 

src/browser/types.ts

@@ -18,7 +18,7 @@ export interface Browser {
     state: Record<string, unknown>;
     applyState: (state: Record<string, unknown>) => void;
     callstackHistory: Callstack;
-    customCommands: string[];
+    customCommands: { name: string; elementScope: boolean }[];
 }
 
 type FunctionProperties<T> = Exclude<

src/runner/browser-env/vite/browser-modules/driver.ts

@@ -27,7 +27,13 @@ export default class ProxyDriver {
         commandWrapper: VoidFunction | undefined,
     ): unknown {
         const monad = webdriverMonad(params, modifier, getWdioPrototype(userPrototype));
-        return monad(window.__testplane__.sessionId, commandWrapper);
+        const browser = monad(window.__testplane__.sessionId, commandWrapper);
+
+        window.__testplane__.customCommands.forEach(({ name, elementScope }) => {
+            browser.addCommand(name, mockCommand(name), elementScope);
+        });
+
+        return browser;
     }
 }
 
@@ -67,25 +73,23 @@ function getAllProtocolCommands(): string[] {
 }
 
 function getMockedProtocolCommands(): PropertiesObject {
-    return [...getAllProtocolCommands(), ...SERVER_HANDLED_COMMANDS, ...window.__testplane__.customCommands].reduce(
-        (acc, commandName) => {
-            acc[commandName] = { value: mockCommand(commandName) };
-            return acc;
-        },
-        {} as PropertiesObject,
-    );
+    return [...getAllProtocolCommands(), ...SERVER_HANDLED_COMMANDS].reduce((acc, commandName) => {
+        acc[commandName] = { value: mockCommand(commandName) };
+        return acc;
+    }, {} as PropertiesObject);
 }
 
 function mockCommand(commandName: string): ProtocolCommandFn {
-    return async (...args: unknown[]): Promise<unknown> => {
+    return async function (this: WebdriverIO.Browser | WebdriverIO.Element, ...args: unknown[]): Promise<unknown> {
         const { socket } = window.__testplane__;
         const timeout = getCommandTimeout(commandName);
+        const element = isWdioElement(this) ? this : undefined;
 
         try {
             // TODO: remove type casting after https://github.com/socketio/socket.io/issues/4925
             const [error, result] = (await socket
                 .timeout(timeout)
-                .emitWithAck(BrowserEventNames.runBrowserCommand, { name: commandName, args })) as [
+                .emitWithAck(BrowserEventNames.runBrowserCommand, { name: commandName, args, element })) as [
                 err: null | Error,
                 result?: unknown,
             ];
@@ -167,3 +171,7 @@ function truncate(value: string, maxLen: number): string {
 
     return `${value.slice(0, maxLen - 3)}...`;
 }
+
+function isWdioElement(ctx: WebdriverIO.Browser | WebdriverIO.Element): ctx is WebdriverIO.Element {
+    return Boolean((ctx as WebdriverIO.Element).elementId);
+}

src/runner/browser-env/vite/browser-modules/mock/@wdio-logger.ts

@@ -0,0 +1,10 @@
+export default function getLogger(): typeof console {
+    return {
+        log: (): void => {},
+        info: (): void => {},
+        warn: (): void => {},
+        error: (): void => {},
+    } as unknown as typeof console;
+}
+
+getLogger.setLogLevelsConfig = (): void => {};

src/runner/browser-env/vite/browser-modules/mock/default-module.ts

@@ -0,0 +1 @@
+export default (): void => {};

src/runner/browser-env/vite/browser-modules/mock/import-meta-resolve.ts

@@ -0,0 +1 @@
+export const resolve = (): void => {};

src/runner/browser-env/vite/browser-modules/types.ts

@@ -19,6 +19,7 @@ export enum BrowserEventNames {
 export interface BrowserRunBrowserCommandPayload {
     name: string;
     args: unknown[];
+    element?: WebdriverIO.Element;
 }
 
 export interface BrowserRunExpectMatcherPayload {
@@ -59,7 +60,7 @@ export interface WorkerInitializePayload {
     sessionId: WebdriverIO.Browser["sessionId"];
     capabilities: WebdriverIO.Browser["capabilities"];
     requestedCapabilities: WebdriverIO.Browser["options"]["capabilities"];
-    customCommands: string[];
+    customCommands: { name: string; elementScope: boolean }[];
     // TODO: use BrowserConfig type after migrate to esm
     config: {
         automationProtocol: "webdriver" | "devtools";

src/runner/browser-env/vite/plugins/generate-index-html.ts

@@ -14,21 +14,16 @@ import type { Plugin, Rollup } from "vite";
 const debug = createDebug("vite:plugin:generateIndexHtml");
 
 // modules that used only in NodeJS environment and don't need to be compiled
-const MODULES_TO_MOCK = ["import-meta-resolve", "puppeteer-core", "archiver", "@wdio/repl"];
+const DEFAULT_MODULES_TO_MOCK = ["puppeteer-core", "archiver", "@wdio/repl"];
 const POLYFILLS = [...builtinModules, ...builtinModules.map(m => `node:${m}`)];
 
 const virtualDriverModuleId = "virtual:@testplane/driver";
-const virtualMockModuleId = "virtual:@testplane/mock";
 
 const virtualModules = {
     driver: {
         id: virtualDriverModuleId,
         resolvedId: `\0${virtualDriverModuleId}`,
     },
-    mock: {
-        id: virtualMockModuleId,
-        resolvedId: `\0${virtualMockModuleId}`,
-    },
 };
 
 export const plugin = async (): Promise<Plugin[]> => {
@@ -46,6 +41,15 @@ export const plugin = async (): Promise<Plugin[]> => {
 
     const automationProtocolPath = `/@fs${driverModulePath}`;
 
+    const mockDefaultModulePath = path.resolve(browserModulesPath, "mock/default-module.js");
+    const mockImportMetaResolvePath = path.resolve(browserModulesPath, "mock/import-meta-resolve.js");
+    const mockWdioLoggerPath = path.resolve(browserModulesPath, "mock/@wdio-logger.js");
+
+    const modulesToMock = DEFAULT_MODULES_TO_MOCK.reduce((acc, val) => _.set(acc, val, mockDefaultModulePath), {
+        "@wdio/logger": mockWdioLoggerPath,
+        "import-meta-resolve": mockImportMetaResolvePath,
+    }) as Record<string, string>;
+
     return [
         {
             name: "testplane:generateIndexHtml",
@@ -114,19 +118,15 @@ export const plugin = async (): Promise<Plugin[]> => {
                     return polyfillPath(id.replace("/promises", ""));
                 }
 
-                if (MODULES_TO_MOCK.includes(id)) {
-                    return virtualModules.mock.resolvedId;
+                if (Object.keys(modulesToMock).includes(id)) {
+                    return modulesToMock[id];
                 }
             },
 
             load: (id: string): Rollup.LoadResult | void => {
                 if (id === virtualModules.driver.resolvedId) {
                     return `export const automationProtocolPath = ${JSON.stringify(automationProtocolPath)};`;
                 }
-
-                if (id === virtualModules.mock.resolvedId) {
-                    return ["export default () => {};", "export const resolve = () => ''"].join("\n");
-                }
             },
 
             transform(code, id): Rollup.TransformResult {

src/worker/browser-env/runner/test-runner/index.ts

@@ -104,16 +104,23 @@ export class TestRunner extends NodejsEnvTestRunner {
         const { publicAPI: session } = browser;
 
         return async (payload, cb): Promise<void> => {
-            const { name, args } = payload;
-            const cmdName = name as keyof typeof session;
-
-            if (typeof session[cmdName] !== "function") {
-                cb([prepareData<Error>(new Error(`"browser.${name}" does not exists in browser instance`))]);
+            const { name, args, element } = payload;
+
+            const wdioInstance = await getWdioInstance(session, element);
+            const wdioInstanceName = element ? "element" : "browser";
+            const cmdName = name as keyof typeof wdioInstance;
+
+            if (typeof wdioInstance[cmdName] !== "function") {
+                cb([
+                    prepareData<Error>(
+                        new Error(`"${wdioInstanceName}.${name}" does not exists in ${wdioInstanceName} instance`),
+                    ),
+                ]);
                 return;
             }
 
             try {
-                const result = await (session[cmdName] as (...args: unknown[]) => Promise<unknown>)(...args);
+                const result = await (wdioInstance[cmdName] as (...args: unknown[]) => Promise<unknown>)(...args);
 
                 if (_.isError(result)) {
                     return cb([prepareData<Error>(result)]);
@@ -209,3 +216,20 @@ function transformExpectArg(arg: any): unknown {
 
     return arg;
 }
+
+async function getWdioInstance(
+    session: WebdriverIO.Browser,
+    element?: WebdriverIO.Element,
+): Promise<WebdriverIO.Browser | WebdriverIO.Element> {
+    const wdioInstance = element ? await session.$(element) : session;
+
+    if (isWdioElement(wdioInstance) && !wdioInstance.selector) {
+        wdioInstance.selector = element?.selector as Selector;
+    }
+
+    return wdioInstance;
+}
+
+function isWdioElement(ctx: WebdriverIO.Browser | WebdriverIO.Element): ctx is WebdriverIO.Element {
+    return Boolean((ctx as WebdriverIO.Element).elementId);
+}

test/src/browser/existing-browser.js

@@ -173,10 +173,12 @@ describe("ExistingBrowser", () => {
                 });
 
                 await initBrowser_(browser);
-                session.addCommand("foo", () => {});
+                session.addCommand("foo", () => {}, false);
+                session.addCommand("foo", () => {}, true);
 
                 assert.isNotEmpty(browser.customCommands);
-                assert.include(browser.customCommands, "foo");
+                assert.deepInclude(browser.customCommands, { name: "foo", elementScope: false });
+                assert.deepInclude(browser.customCommands, { name: "foo", elementScope: true });
             });
         });
 

test/src/browser/history/commands.js

@@ -11,6 +11,7 @@ describe("commands-history", () => {
                     "$",
                     "action",
                     "actions",
+                    "addCommand",
                     "call",
                     "custom$$",
                     "custom$",
@@ -26,6 +27,7 @@ describe("commands-history", () => {
                     "mockClearAll",
                     "mockRestoreAll",
                     "newWindow",
+                    "overwriteCommand",
                     "pause",
                     "react$$",
                     "react$",
@@ -61,8 +63,8 @@ describe("commands-history", () => {
                     "dragAndDrop",
                     "getAttribute",
                     "getCSSProperty",
-                    "getComputedRole",
                     "getComputedLabel",
+                    "getComputedRole",
                     "getHTML",
                     "getLocation",
                     "getProperty",