Make ReactWrapper and ShallowWrapper iterable

enzymejs · Sep 12, 2016 · d36dc6c · d36dc6c
1 parent 69b7a18
commit d36dc6c
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 0 deletions.
diff --git a/src/ReactWrapper.jsx b/src/ReactWrapper.jsx
@@ -28,6 +28,7 @@ import {
   propsOfNode,
   typeOfNode,
   displayNameOfNode,
+  ITERATOR_SYMBOL,
 } from './Utils';
 import {
   debugInsts,
@@ -106,6 +107,26 @@ export default class ReactWrapper {
     );
   }
 
+  /**
+   * Makes a wrapper iterable, which is useful when using destructive
+   * assignment with `find()`
+   * @example
+   * const wrapper = shallow(<MyComponent />)
+   * const [fistLink, secondLink] = wrapper.find('a')
+   * @returns {Object}
+   */
+  [ITERATOR_SYMBOL]() {
+    let index = 0;
+    return {
+      next: () => {
+        if (index >= this.nodes.length) {
+          return { done: true };
+        }
+        return { done: false, value: this.nodes[index++] };
+      },
+    };
+  }
+
   /**
    * If the root component contained a ref, you can access it here
    * and get a wrapper around it.

diff --git a/src/ShallowWrapper.js b/src/ShallowWrapper.js
@@ -15,6 +15,7 @@ import {
   isReactElementAlike,
   displayNameOfNode,
   isFunctionalComponent,
+  ITERATOR_SYMBOL,
 } from './Utils';
 import {
   debugNodes,
@@ -101,6 +102,26 @@ export default class ShallowWrapper {
     this.complexSelector = new ComplexSelector(buildPredicate, findWhereUnwrapped, childrenOfNode);
   }
 
+  /**
+   * Makes a wrapper iterable, which is useful when using destructive
+   * assignment with `find()`
+   * @example
+   * const wrapper = shallow(<MyComponent />)
+   * const [fistLink, secondLink] = wrapper.find('a')
+   * @returns {Object}
+   */
+  [ITERATOR_SYMBOL]() {
+    let index = 0;
+    return {
+      next: () => {
+        if (index >= this.nodes.length) {
+          return { done: true };
+        }
+        return { done: false, value: this.nodes[index++] };
+      },
+    };
+  }
+
   /**
    * Gets the instance of the component being rendered as the root node passed into `shallow()`.
    *

diff --git a/src/Utils.js b/src/Utils.js
@@ -12,6 +12,10 @@ import {
   REACT15,
 } from './version';
 
+export const ITERATOR_SYMBOL = (
+  typeof Symbol === 'function' && Symbol.iterator
+) || '@@iterator';
+
 function internalInstanceKey(node) {
   return Object.keys(Object(node)).filter(key => key.match(/^__reactInternalInstance\$/))[0];
 }

diff --git a/test/ReactWrapper-spec.jsx b/test/ReactWrapper-spec.jsx
@@ -3019,4 +3019,31 @@ describeWithDOM('mount', () => {
       });
     });
   });
+
+  describe('@@iterator', () => {
+    it('should be iterable', () => {
+      class Foo extends React.Component {
+        render() {
+          return (
+            <div>
+              <a href="#1">Hello</a>
+              <a href="#2">Hello</a>
+              <a href="#3">Hello</a>
+              <a href="#4">Hello</a>
+            </div>
+          );
+        }
+      }
+      const wrapper = mount(<Foo />);
+      const [a, b, c, d] = wrapper.find('a');
+      const a1 = wrapper.find('a').get(0);
+      const b1 = wrapper.find('a').get(1);
+      const c1 = wrapper.find('a').get(2);
+      const d1 = wrapper.find('a').get(3);
+      expect(a1).to.equal(a);
+      expect(b1).to.equal(b);
+      expect(c1).to.equal(c);
+      expect(d1).to.equal(d);
+    });
+  });
 });
diff --git a/test/ShallowWrapper-spec.jsx b/test/ShallowWrapper-spec.jsx
@@ -3575,4 +3575,31 @@ describe('shallow', () => {
     });
   });
 
+  describe('@@iterator', () => {
+    it('should be iterable', () => {
+      class Foo extends React.Component {
+        render() {
+          return (
+            <div>
+              <a href="#1">Hello</a>
+              <a href="#2">Hello</a>
+              <a href="#3">Hello</a>
+              <a href="#4">Hello</a>
+            </div>
+          );
+        }
+      }
+      const wrapper = shallow(<Foo />);
+      const [a, b, c, d] = wrapper.find('a');
+      const a1 = wrapper.find('a').get(0);
+      const b1 = wrapper.find('a').get(1);
+      const c1 = wrapper.find('a').get(2);
+      const d1 = wrapper.find('a').get(3);
+      expect(a1).to.equal(a);
+      expect(b1).to.equal(b);
+      expect(c1).to.equal(c);
+      expect(d1).to.equal(d);
+    });
+  });
+
 });