fix: make chainable methods work right

jcoreio · Dec 18, 2020 · 56d2795 · 56d2795
1 parent e778c7f
commit 56d2795
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 10 deletions.
diff --git a/src/index.js b/src/index.js
@@ -82,16 +82,52 @@ module.exports = (chai, utils) => {
     return name !== '_obj' && typeof propertyDescs[name].get === 'function'
   })
 
+  const isChainableMethod = (getterName) =>
+    Object.prototype.hasOwnProperty.call(
+      Assertion.prototype.__methods,
+      getterName
+    )
+
   getterNames.forEach((getterName) => {
-    Object.defineProperty(WaitFor.prototype, getterName, {
-      get() {
-        return new WaitFor(this.options, () => {
-          const assertion = this.buildAssertion()
-          return assertion[getterName]
-        })
-      },
-      configurable: true,
-    })
+    if (isChainableMethod(getterName)) {
+      Object.defineProperty(WaitFor.prototype, getterName, {
+        get() {
+          const obj = new WaitFor(this.options, () => {
+            const assertion = this.buildAssertion()
+            return assertion[getterName]
+          })
+          function chainable() {
+            return new WaitFor(this.options, () => {
+              const assertion = this.buildAssertion()
+              return assertion[getterName].apply(assertion, arguments)
+            })
+          }
+          for (const methodName of methodNames) {
+            chainable[methodName] = obj[methodName].bind(obj)
+          }
+          for (const getterName of getterNames) {
+            Object.defineProperty(chainable, getterName, {
+              get() {
+                return obj[getterName]
+              },
+              configurable: true,
+            })
+          }
+          return chainable
+        },
+        configurable: true,
+      })
+    } else {
+      Object.defineProperty(WaitFor.prototype, getterName, {
+        get() {
+          return new WaitFor(this.options, () => {
+            const assertion = this.buildAssertion()
+            return assertion[getterName]
+          })
+        },
+        configurable: true,
+      })
+    }
   })
 }
 

diff --git a/test/index.js b/test/index.js
@@ -45,7 +45,33 @@ describe('waitFor', function () {
     ])
     expect(i).to.equal(3)
   })
-
+  it('.to.include works', async function () {
+    let i = 0
+    const values = [
+      { foo: 1, bar: 1 },
+      { foo: 2, bar: 1 },
+      { foo: 3, bar: 1 },
+      { foo: 4, bar: 1 },
+    ]
+    await Promise.all([
+      waitFor(() => values[i++]).to.include({ foo: 3 }),
+      clock.tickAsync(501),
+    ])
+    expect(i).to.equal(3)
+  })
+  it('.to.include.all.keys works', async function () {
+    let i = 0
+    let values = {}
+    await Promise.all([
+      waitFor(() => (values = { ...values, [i++]: true })).to.include.all.keys(
+        '1',
+        '2',
+        '3'
+      ),
+      clock.tickAsync(501),
+    ])
+    expect(i).to.equal(4)
+  })
   it('works when an assertion takes longer than retryInterval', async function () {
     let i = 0
     await Promise.all([