From aac57b30e8127a59c8ccc02cb5c50a87972f68c9 Mon Sep 17 00:00:00 2001
From: Andrey Lunyov <alunyov@meta.com>
Date: Wed, 3 Jan 2024 21:32:50 -0800
Subject: [PATCH] useFragmentInternal should suspend if environment changes and
 there is a query in flight

Summary:
There is a bug in the new hooks implementation, where the `useFragmentInternal` would not suspend after the change of the environment (while the fragment selector are the same) with the active operation in flight.

This diff fixes that by checking if the environment has changed and we would throw the promise for pending operation in flight.

Reviewed By: voideanvalue

Differential Revision: D52516772

fbshipit-source-id: 13a61114e6b7718af548cf45d2e4685bcd121993
---
 .../__tests__/useFragmentNode-test.js         | 27 +++++++++++++++++++
 .../useFragmentInternal_EXPERIMENTAL.js       |  1 +
 2 files changed, 28 insertions(+)

diff --git a/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js b/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js
index 5994469f365d3..ef65e9a1fecac 100644
--- a/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js
+++ b/packages/react-relay/relay-hooks/__tests__/useFragmentNode-test.js
@@ -661,6 +661,33 @@ describe.each([
       ]);
     });
 
+    it('should supsend when the environment changes and there is query in flight', () => {
+      const renderer = renderSingularFragment();
+      assertFragmentResults([
+        {
+          data: {
+            id: '1',
+            name: 'Alice',
+            profile_picture: null,
+            ...createFragmentRef('1', singularQuery),
+          },
+        },
+      ]);
+
+      const newEnvironment = createMockEnvironment();
+
+      internalAct(() => {
+        // Let there be an operation in flight
+        fetchQuery(newEnvironment, singularQuery).subscribe({});
+
+        setEnvironment(newEnvironment);
+      });
+
+      // It should suspend when the environment changes and there is a query
+      // in flight.
+      expect(renderer.toJSON()).toEqual('Singular Fallback');
+    });
+
     it('should re-read and resubscribe to fragment when fragment pointers change', () => {
       renderSingularFragment();
       assertRenderBatch([
diff --git a/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js b/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js
index 08f6407b3200d..9318dcab05c93 100644
--- a/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js
+++ b/packages/react-relay/relay-hooks/experimental/useFragmentInternal_EXPERIMENTAL.js
@@ -505,6 +505,7 @@ function useFragmentInternal_EXPERIMENTAL(
     // We only suspend when the component is first trying to mount or changing
     // selectors, not if data becomes missing later:
     if (
+      environment !== previousEnvironment ||
       !committedFragmentSelectorRef.current ||
       !areEqualSelectors(committedFragmentSelectorRef.current, fragmentSelector)
     ) {