Support iterables generically

enzymejs · Oct 27, 2017 · 321c807 · 321c807
1 parent d432c69
commit 321c807
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 15 deletions.
diff --git a/packages/enzyme-adapter-react-13/src/ReactThirteenElementToTree.js b/packages/enzyme-adapter-react-13/src/ReactThirteenElementToTree.js
@@ -1,4 +1,5 @@
 import flatten from 'lodash/flatten';
+import { isIterable } from 'enzyme-adapter-utils';
 
 export function nodeTypeFromType(type) {
   if (typeof type === 'string') {
@@ -23,8 +24,8 @@ export default function elementToTree(el) {
   let rendered = null;
   if (Array.isArray(children)) {
     rendered = flatten(children, true).map(elementToTree);
-  } else if (children && typeof children.toArray === 'function') {
-    rendered = flatten(children.toArray(), true).map(elementToTree);
+  } else if (isIterable(children) && typeof children !== 'string') {
+    rendered = flatten([...children], true).map(elementToTree);
   } else if (typeof children !== 'undefined') {
     rendered = elementToTree(children);
   }

diff --git a/packages/enzyme-adapter-utils/src/Utils.js b/packages/enzyme-adapter-utils/src/Utils.js
@@ -95,6 +95,15 @@ export function nodeTypeFromType(type) {
   return 'function';
 }
 
+export function isIterable(obj) {
+  return (
+    obj != null &&
+    typeof Symbol === 'function' &&
+    typeof Symbol.iterator === 'symbol' &&
+    typeof obj[Symbol.iterator] === 'function'
+  );
+}
+
 export function elementToTree(el) {
   if (el === null || typeof el !== 'object' || !('type' in el)) {
     return el;
@@ -104,8 +113,8 @@ export function elementToTree(el) {
   let rendered = null;
   if (Array.isArray(children)) {
     rendered = flatten(children, true).map(elementToTree);
-  } else if (children && typeof children.toArray === 'function') {
-    rendered = flatten(children.toArray(), true).map(elementToTree);
+  } else if (isIterable(children) && typeof children !== 'string') {
+    rendered = flatten([...children], true).map(elementToTree);
   } else if (typeof children !== 'undefined') {
     rendered = elementToTree(children);
   }

diff --git a/packages/enzyme-test-suite/test/RSTTraversal-spec.jsx b/packages/enzyme-test-suite/test/RSTTraversal-spec.jsx
@@ -189,33 +189,55 @@ describe('RSTTraversal', () => {
         <div key="a" />,
         <div key="b" />,
       ];
+      const divA = $(<div key="a" />);
+      const divB = $(<div key="b" />);
       const node = $(
         <div>
           {twoDivArray}
         </div>,
       );
       treeForEach(node, spy);
       expect(spy.callCount).to.equal(3);
+      const nodes = spy.args.map(arg => arg[0]);
+      expect(nodes).to.deep.equal([node, divA, divB]);
     });
 
-    it('should handle Immutable.js children', () => {
+    it('should handle Map children', () => {
       const spy = sinon.spy();
-      // This object mimics only the toArray functionality of Immutable.js
-      const twoDivImmutable = {
-        toArray() {
-          return [
-            <div key="a" />,
-            <div key="b" />,
-          ];
-        },
-      };
+      const twoDivMap = new Map([
+        [<div key="a" />],
+        [<div key="b" />],
+      ]);
+      const divA = $(<div key="a" />);
+      const divB = $(<div key="b" />);
       const node = $(
         <div>
-          {twoDivImmutable}
+          {twoDivMap}
         </div>,
       );
       treeForEach(node, spy);
       expect(spy.callCount).to.equal(3);
+      const nodes = spy.args.map(arg => arg[0]);
+      expect(nodes).to.deep.equal([node, divA, divB]);
+    });
+
+    it('should handle Set children', () => {
+      const spy = sinon.spy();
+      const twoDivSet = new Set([
+        <div key="a" />,
+        <div key="b" />,
+      ]);
+      const divA = $(<div key="a" />);
+      const divB = $(<div key="b" />);
+      const node = $(
+        <div>
+          {twoDivSet}
+        </div>,
+      );
+      treeForEach(node, spy);
+      expect(spy.callCount).to.equal(3);
+      const nodes = spy.args.map(arg => arg[0]);
+      expect(nodes).to.deep.equal([node, divA, divB]);
     });
 
     it('should not get trapped from empty strings', () => {