Conditionally define iterator method

src/ReactWrapper.jsx

@@ -63,7 +63,7 @@ function filterWhereUnwrapped(wrapper, predicate) {
 /**
  * @class ReactWrapper
  */
-export default class ReactWrapper {
+class ReactWrapper {
 
   constructor(nodes, root, options = {}) {
     if (!global.window && !global.document) {
@@ -107,26 +107,6 @@ 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.
@@ -929,3 +909,12 @@ export default class ReactWrapper {
     unmountComponentAtNode(this.options.attachTo);
   }
 }
+
+
+if (ITERATOR_SYMBOL) {
+  ReactWrapper.prototype[ITERATOR_SYMBOL] = function iterator() {
+    return this.nodes[ITERATOR_SYMBOL]();
+  };
+}
+
+export default ReactWrapper;

src/ShallowWrapper.js

@@ -64,7 +64,7 @@ function filterWhereUnwrapped(wrapper, predicate) {
 /**
  * @class ShallowWrapper
  */
-export default class ShallowWrapper {
+class ShallowWrapper {
 
   constructor(nodes, root, options = {}) {
     if (!root) {
@@ -104,26 +104,6 @@ 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()`.
    *
@@ -1003,3 +983,11 @@ export default class ShallowWrapper {
     });
   }
 }
+
+if (ITERATOR_SYMBOL) {
+  ShallowWrapper.prototype[ITERATOR_SYMBOL] = function iterator() {
+    return this.nodes[ITERATOR_SYMBOL]();
+  };
+}
+
+export default ShallowWrapper;

src/Utils.js

@@ -12,9 +12,7 @@ import {
   REACT15,
 } from './version';
 
-export const ITERATOR_SYMBOL = (
-  typeof Symbol === 'function' && Symbol.iterator
-) || '@@iterator';
+export const ITERATOR_SYMBOL = typeof Symbol === 'function' && Symbol.iterator;
 
 function internalInstanceKey(node) {
   return Object.keys(Object(node)).filter(key => key.match(/^__reactInternalInstance\$/))[0];