From dd18ca42ff5b65379514e9859b4e3b547b0e38bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc-Andr=C3=A9=20Giroux?= <mgiroux0@gmail.com>
Date: Fri, 24 Jun 2016 17:55:11 -0700
Subject: [PATCH] Support route conditional fragments returning an array

Summary:
Gave #896 a try.

modifies `createNode` and `getChildren` to possibly return an array of nodes when a route conditional fragment reference returns an array.

Let me know what you think, It's seems weird to have `createNode` possibly return an Array but I'm not sure how it could be done without it!

[Updates to original PR from wincent: switched `var` to `let`/`const`, avoided an array allocation in `getChildren`.]

Closes https://github.com/facebook/relay/pull/915

Reviewed By: josephsavona

Differential Revision: D3484897

Pulled By: wincent

fbshipit-source-id: 0d101bf63e1a7e5669fc4bd9094e0d35c17c8c40
---
 src/query/RelayQuery.js                | 30 ++++++++-----
 src/query/__tests__/RelayQuery-test.js | 59 +++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 12 deletions(-)

diff --git a/src/query/RelayQuery.js b/src/query/RelayQuery.js
index 79881d0dd4ccd..16c4ea62cb1c7 100644
--- a/src/query/RelayQuery.js
+++ b/src/query/RelayQuery.js
@@ -215,13 +215,19 @@ class RelayQueryNode {
           if (concreteChild == null) {
             return;
           }
-          const node = createNode(
+          const nodeOrNodes = createNode(
             concreteChild,
             this.__route__,
             this.__variables__
           );
-          if (node && node.isIncluded()) {
-            nextChildren.push(node);
+          if (Array.isArray(nodeOrNodes)) {
+            nodeOrNodes.forEach(node => {
+              if (node && node.isIncluded()) {
+                nextChildren.push(node);
+              }
+            });
+          } else if (nodeOrNodes && nodeOrNodes.isIncluded()) {
+            nextChildren.push(nodeOrNodes);
           }
         });
       }
@@ -1328,7 +1334,7 @@ function createNode(
   concreteNode: mixed,
   route: RelayMetaRoute,
   variables: Variables
-): ?RelayQueryNode {
+): ?RelayQueryNode | Array<?RelayQueryNode> {
   invariant(
     typeof concreteNode === 'object' &&
     concreteNode !== null,
@@ -1366,13 +1372,15 @@ function createNode(
     type = RelayQuerySubscription;
   } else if (concreteNode instanceof RelayRouteFragment) {
     const fragment = concreteNode.getFragmentForRoute(route);
-    if (fragment) {
-      // may be null if no value was defined for this route.
-      return createNode(
-        fragment,
-        route,
-        variables
-      );
+    // May be null if no value was defined for this route.
+    if (Array.isArray(fragment)) {
+      // A route-conditional function may return a single fragment reference
+      // or an array of fragment references.
+      return fragment.map(frag => {
+        return createNode(frag, route, variables);
+      });
+    } else if (fragment) {
+      return createNode(fragment, route, variables);
     }
     return null;
   } else if (concreteNode instanceof RelayFragmentReference) {
diff --git a/src/query/__tests__/RelayQuery-test.js b/src/query/__tests__/RelayQuery-test.js
index 8a44f9f114353..f40c9b49d1672 100644
--- a/src/query/__tests__/RelayQuery-test.js
+++ b/src/query/__tests__/RelayQuery-test.js
@@ -292,7 +292,7 @@ describe('RelayQuery', () => {
           type: 'Node',
         });
         expect(field.canHaveSubselections()).toBe(true);
-        var children = field.getChildren();
+        const children = field.getChildren();
         expect(children.length).toBe(2);
         expect(children[0]).toBe(child);
         expect(children[1]).toBe(fragment);
@@ -486,5 +486,62 @@ describe('RelayQuery', () => {
         {name: 'size', value: 'override'},
       ]);
     });
+
+    it('expands route conditional fragments', () => {
+      const innerFragment1 = Relay.QL`
+        fragment on User {
+          id,
+          profilePicture(size:$size) {
+            uri,
+          },
+        }
+      `;
+      const innerFragment2 = Relay.QL`
+        fragment on User {
+          id,
+          firstName
+        }
+      `;
+      const reference1 = new RelayFragmentReference(
+        () => innerFragment1,
+        {
+          size: 'default',
+        },
+        {
+          size: QueryBuilder.createCallVariable('outerSize'),
+        }
+      );
+      const reference2 = new RelayFragmentReference(() => innerFragment2, {}, {});
+      const fragment = getNode(Relay.QL`
+        fragment on User {
+          id,
+          ${route => reference1},
+          ${route => [reference2]}
+        }
+      `, {
+        outerSize: 'override',
+      });
+
+      const children = fragment.getChildren();
+      expect(children.length).toBe(3);
+      expect(children[0].getSchemaName()).toBe('id');
+
+      expect(children[1] instanceof RelayQuery.Fragment);
+      expect(children[1].getType()).toBe('User');
+      let grandchildren = children[1].getChildren();
+      expect(grandchildren.length).toBe(2);
+      expect(grandchildren[0].getSchemaName()).toBe('id');
+      expect(grandchildren[1].getSchemaName()).toBe('profilePicture');
+      expect(grandchildren[1].getCallsWithValues()).toEqual([
+        {name: 'size', value: 'override'},
+      ]);
+
+      expect(children[2] instanceof RelayQuery.Fragment);
+      expect(children[2].getType()).toBe('User');
+      grandchildren = children[2].getChildren();
+      expect(grandchildren.length).toBe(2);
+      expect(grandchildren[0].getSchemaName()).toBe('id');
+      expect(grandchildren[1].getSchemaName()).toBe('firstName');
+    });
   });
 });