From 4eb13770a70417a4a7d9372d622ec1916e59df84 Mon Sep 17 00:00:00 2001
From: Jay Chen <jchencareer@gmail.com>
Date: Mon, 21 Mar 2022 11:38:17 +0100
Subject: [PATCH] allow chain custom commands

---
 packages/wdio-utils/src/shim.ts              |  4 ++--
 packages/wdio-utils/tests/shim-async.test.ts | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/packages/wdio-utils/src/shim.ts b/packages/wdio-utils/src/shim.ts
index a78ee8aa99a..5e7eb15057c 100644
--- 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
index 32cd491905a..706d196a01b 100644
--- 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()