From f001a5eea9ddb2a04609959aeba733618fa6f954 Mon Sep 17 00:00:00 2001
From: Lee Byron <lee.byron@robinhood.com>
Date: Tue, 2 Apr 2019 15:30:12 -0700
Subject: [PATCH] Fix case where loop callback returns something other than
 element, improve tests, ensure key is always provided as a string

---
 react-loops.js      | 11 ++++-------
 react-loops.test.js | 44 +++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/react-loops.js b/react-loops.js
index 6d1cc78..83a1ceb 100644
--- a/react-loops.js
+++ b/react-loops.js
@@ -132,7 +132,7 @@ function For(props) {
 }
 
 function mapIteration(mapper, item, index, length, key) {
-  var itemChildren =
+  var result =
     mapper.length === 1
       ? mapper(item)
       : mapper(item, {
@@ -142,13 +142,10 @@ function mapIteration(mapper, item, index, length, key) {
           isFirst: index === 0,
           isLast: index === length - 1
         });
-  if (React.Children.count(itemChildren) === 1) {
-    var child = React.Children.only(itemChildren);
-    return child.props.hasOwnProperty("key")
-      ? child
-      : React.cloneElement(child, { key: key });
+  if (React.isValidElement(result) && !result.props.hasOwnProperty("key")) {
+    return React.cloneElement(result, { key: String(key) });
   }
-  return itemChildren;
+  return result;
 }
 
 /**
diff --git a/react-loops.test.js b/react-loops.test.js
index 4ee2805..74bb644 100644
--- a/react-loops.test.js
+++ b/react-loops.test.js
@@ -3,9 +3,17 @@ import TestRenderer from "react-test-renderer";
 import { For, If, ElseIf, Else } from "./react-loops.js";
 
 function expectRenderToEqual(actual, expected) {
-  expect(TestRenderer.create(actual).toJSON()).toEqual(
-    TestRenderer.create(expected).toJSON()
-  );
+  const consoleError = console.error;
+  try {
+    console.error = message => {
+      throw new Error(message);
+    };
+    expect(TestRenderer.create(actual).toJSON()).toEqual(
+      TestRenderer.create(expected).toJSON()
+    );
+  } finally {
+    console.error = consoleError;
+  }
 }
 
 function expectRenderToThrow(actual, error) {
@@ -90,6 +98,36 @@ describe("react-loops", () => {
         );
       });
 
+      it("loops with a non-element string return", () => {
+        const list = ["A", "B", "C"];
+        expectRenderToEqual(
+          <For of={list} as={item => item} />,
+          <>
+            {"A"}
+            {"B"}
+            {"C"}
+          </>
+        );
+      });
+
+      it("loops with a non-element Array return", () => {
+        const list = ["A", "B", "C"];
+        expectRenderToEqual(
+          <For
+            of={list}
+            as={item => [<li key="0">{item}</li>, <li key="1">{item}</li>]}
+          />,
+          <>
+            <li key="woof">A</li>
+            <li>A</li>
+            <li>B</li>
+            <li>B</li>
+            <li>C</li>
+            <li>C</li>
+          </>
+        );
+      });
+
       it("loops an Array with metadata", () => {
         const list = ["A", "B", "C"];
         expectRenderToEqual(