Merge pull request enzymejs#594 from Aweary/iterable-wrapper

Make ReactWrapper and ShallowWrapper iterable
joeduncan · Oct 11, 2016 · 746eac4 · 746eac4
2 parents d041930 + 5b9022a
commit 746eac4
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 2 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,
@@ -62,7 +63,7 @@ function filterWhereUnwrapped(wrapper, predicate) {
 /**
  * @class ReactWrapper
  */
-export default class ReactWrapper {
+class ReactWrapper {
 
   constructor(nodes, root, options = {}) {
     if (!global.window && !global.document) {
@@ -908,3 +909,15 @@ export default class ReactWrapper {
     unmountComponentAtNode(this.options.attachTo);
   }
 }
+
+
+if (ITERATOR_SYMBOL) {
+  Object.defineProperty(ReactWrapper.prototype, ITERATOR_SYMBOL, {
+    configurable: true,
+    value: function iterator() {
+      return this.nodes[ITERATOR_SYMBOL]();
+    },
+  });
+}
+
+export default ReactWrapper;
diff --git a/src/ShallowWrapper.js b/src/ShallowWrapper.js
@@ -16,6 +16,7 @@ import {
   displayNameOfNode,
   isFunctionalComponent,
   isCustomComponentElement,
+  ITERATOR_SYMBOL,
 } from './Utils';
 import {
   debugNodes,
@@ -63,7 +64,7 @@ function filterWhereUnwrapped(wrapper, predicate) {
 /**
  * @class ShallowWrapper
  */
-export default class ShallowWrapper {
+class ShallowWrapper {
 
   constructor(nodes, root, options = {}) {
     if (!root) {
@@ -982,3 +983,14 @@ export default class ShallowWrapper {
     });
   }
 }
+
+if (ITERATOR_SYMBOL) {
+  Object.defineProperty(ShallowWrapper.prototype, ITERATOR_SYMBOL, {
+    configurable: true,
+    value: function iterator() {
+      return this.nodes[ITERATOR_SYMBOL]();
+    },
+  });
+}
+
+export default ShallowWrapper;
diff --git a/src/Utils.js b/src/Utils.js
@@ -12,6 +12,8 @@ import {
   REACT15,
 } from './version';
 
+export const ITERATOR_SYMBOL = typeof Symbol === 'function' && Symbol.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
@@ -16,6 +16,7 @@ import {
   render,
   ReactWrapper,
 } from '../src';
+import { ITERATOR_SYMBOL } from '../src/Utils';
 import { REACT013, REACT014, REACT15 } from '../src/version';
 
 describeWithDOM('mount', () => {
@@ -3024,4 +3025,31 @@ describeWithDOM('mount', () => {
       });
     });
   });
+
+  describeIf(ITERATOR_SYMBOL, '@@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
@@ -4,6 +4,7 @@ import sinon from 'sinon';
 
 import { shallow, render, ShallowWrapper } from '../src/';
 import { describeIf, itIf, itWithData, generateEmptyRenderData } from './_helpers';
+import { ITERATOR_SYMBOL } from '../src/Utils';
 import { REACT013, REACT15 } from '../src/version';
 
 describe('shallow', () => {
@@ -3626,4 +3627,31 @@ describe('shallow', () => {
       expect(underwater.is(RendersDOM)).to.equal(true);
     });
   });
+
+  describeIf(ITERATOR_SYMBOL, '@@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);
+    });
+  });
 });