allow chain custom commands

webdriverio · Mar 21, 2022 · 4eb1377 · 4eb1377
1 parent 8ee4326
commit 4eb1377
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/packages/wdio-utils/src/shim.ts b/packages/wdio-utils/src/shim.ts
@@ -251,7 +251,7 @@ let wrapCommand = function wrapCommand<T>(commandName: string, fn: Function): (.
                      * await $('foo').$('bar')
                      * ```
                      */
-                    if (ELEMENT_QUERY_COMMANDS.includes(prop)) {
+                    if (ELEMENT_QUERY_COMMANDS.includes(prop) || prop.endsWith('$')) {
                         // this: WebdriverIO.Element
                         return wrapCommand(prop, function (this: any, ...args: any) {
                             return this[prop].apply(this, args)
@@ -336,7 +336,7 @@ let wrapCommand = function wrapCommand<T>(commandName: string, fn: Function): (.
          */
         const command = hasWdioSyncSupport && wdioSync && Boolean(global.browser) && !runAsync && !asyncSpec
             ? wdioSync!.wrapCommand(commandName, fn)
-            : ELEMENT_QUERY_COMMANDS.includes(commandName)
+            : ELEMENT_QUERY_COMMANDS.includes(commandName) || commandName.endsWith('$')
                 ? chainElementQuery
                 : wrapCommandFn
 

diff --git a/packages/wdio-utils/tests/shim-async.test.ts b/packages/wdio-utils/tests/shim-async.test.ts
@@ -246,6 +246,25 @@ describe('wrapCommand', () => {
         expect(scope.getTagName).toBeCalledTimes(1)
     })
 
+    it('allows to chain element promises for custom command', async () => {
+        const rawCommand = jest.fn()
+        const scope: Partial<BrowserObject> = {
+            options: {
+                beforeCommand: jest.fn(),
+                afterCommand: jest.fn()
+            },
+            getTagName: jest.fn().mockResolvedValue('Yayy'),
+            user$: rawCommand
+        }
+        rawCommand.mockReturnValue(Promise.resolve(scope))
+        const commandB = wrapCommand('user$', rawCommand)
+        expect(await commandB.call(scope, 'bar').user$('foo').getTagName()).toBe('Yayy')
+        expect(scope.user$).toBeCalledTimes(2)
+        expect(scope.user$).toBeCalledWith('bar')
+        expect(scope.user$).toBeCalledWith('foo')
+        expect(scope.getTagName).toBeCalledTimes(1)
+    })
+
     it('allows to access indexed element', async () => {
         const rawCommand$ = jest.fn()
         const rawCommand$$ = jest.fn()