From 3036f2e0df54e49fcac0e4074ba42cd49b6ead7c Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Fri, 2 Dec 2022 18:41:02 -0500
Subject: [PATCH] Fix transition state tests

---
 integration/transition-state-test.ts | 53 ++++++++++++++++++----------
 packages/remix-react/components.tsx  | 50 +++++++++++++-------------
 2 files changed, 60 insertions(+), 43 deletions(-)

diff --git a/integration/transition-state-test.ts b/integration/transition-state-test.ts
index 4d973cb2a83..4efec22f3d8 100644
--- a/integration/transition-state-test.ts
+++ b/integration/transition-state-test.ts
@@ -253,11 +253,14 @@ test.describe("rendering", () => {
           pathname: "/",
           search: "?redirected",
           hash: "",
-          state: {
-            isRedirect: true,
-            setCookie: false,
-            type: "loader",
-          },
+          // These were private API for transition manager that are no longer
+          // needed with the new router so OK to disappear
+          // state: {
+          //   isRedirect: true,
+          //   setCookie: false,
+          //   type: "loader",
+          // },
+          state: null,
           key: expect.any(String),
         },
       },
@@ -288,6 +291,9 @@ test.describe("rendering", () => {
           state: null,
           key: expect.any(String),
         },
+        // TODO This fails because we don't expose the "submission" on loader
+        // submissions from the new router, but we think we probably should
+        // even though it's a "loading" navigation
         submission: {
           action: `/${STATES.SUBMITTING_LOADER}`,
           encType: "application/x-www-form-urlencoded",
@@ -338,11 +344,14 @@ test.describe("rendering", () => {
           pathname: "/",
           search: "?redirected",
           hash: "",
-          state: {
-            isRedirect: true,
-            setCookie: false,
-            type: "loaderSubmission",
-          },
+          // These were private API for transition manager that are no longer
+          // needed with the new router so OK to disappear
+          // state: {
+          //   isRedirect: true,
+          //   setCookie: false,
+          //   type: "loaderSubmission",
+          // },
+          state: null,
           key: expect.any(String),
         },
         submission: {
@@ -448,11 +457,14 @@ test.describe("rendering", () => {
           pathname: "/",
           search: "?redirected",
           hash: "",
-          state: {
-            isRedirect: true,
-            setCookie: false,
-            type: "action",
-          },
+          // These were private API for transition manager that are no longer
+          // needed with the new router so OK to disappear
+          // state: {
+          //   isRedirect: true,
+          //   setCookie: false,
+          //   type: "action",
+          // },
+          state: null,
           key: expect.any(String),
         },
         submission: {
@@ -487,10 +499,15 @@ test.describe("rendering", () => {
           pathname: "/",
           search: "?redirected",
           hash: "",
+          // These were private API for transition manager that are no longer
+          // needed with the new router so OK to disappear
+          // state: {
+          //   isRedirect: true,
+          //   setCookie: false,
+          //   type: "fetchAction",
+          // },
           state: {
-            isRedirect: true,
-            setCookie: false,
-            type: "fetchAction",
+            isFetchActionRedirect: true,
           },
           key: expect.any(String),
         },
diff --git a/packages/remix-react/components.tsx b/packages/remix-react/components.tsx
index 0b559ac9db0..ca01bcc1564 100644
--- a/packages/remix-react/components.tsx
+++ b/packages/remix-react/components.tsx
@@ -8,6 +8,7 @@ import type {
   AgnosticDataRouteMatch,
   AgnosticDataRouteObject,
   ErrorResponse,
+  Navigation,
 } from "@remix-run/router";
 import type {
   LinkProps,
@@ -71,6 +72,7 @@ import type {
   TransitionStates,
 } from "./transition";
 import { IDLE_TRANSITION, IDLE_FETCHER } from "./transition";
+import { labeledStatement } from "@babel/types";
 
 function useDataRouterContext() {
   let context = React.useContext(DataRouterContext);
@@ -945,7 +947,11 @@ export function useActionData<T = AppData>(): SerializeFrom<T> | undefined {
  * @see https://remix.run/api/remix#usetransition
  */
 export function useTransition(): Transition {
+  let currentLocation = useLocation();
   let navigation = useNavigation();
+  let lastNavigationRef = React.useRef<Navigation | null>(null);
+  let lastNavigation = lastNavigationRef.current;
+  lastNavigationRef.current = navigation;
 
   // TODO: Should we populate navigation.formData on <Form method="get"> even
   // though we've already move the data onto URLSearchParams.
@@ -1071,31 +1077,25 @@ export function useTransition(): Transition {
           return transition;
         }
       }
-    } else {
-      // TODO: How can we detect this?
-      let wasNormalRedirect = false;
-      if (wasNormalRedirect) {
-        let transition: TransitionStates["LoadingRedirect"] = {
-          location,
-          state,
-          submission: undefined,
-          type: "normalRedirect",
-        };
-        return transition;
-      }
-
-      // TODO: How can we detect a fetch action redirect???  Do we need to
-      // check useFetchers?  Or could we somehow look at location key?
-      let wasFetchActionRedirect = false;
-      if (wasFetchActionRedirect) {
-        let transition: TransitionStates["LoadingFetchActionRedirect"] = {
-          location,
-          state,
-          submission: undefined,
-          type: "fetchActionRedirect",
-        };
-        return transition;
-      }
+    } else if (
+      lastNavigation?.state === "loading" &&
+      lastNavigation.location.key !== navigation.location?.key
+    ) {
+      let transition: TransitionStates["LoadingRedirect"] = {
+        location,
+        state,
+        submission: undefined,
+        type: "normalRedirect",
+      };
+      return transition;
+    } else if (location.state?.isFetchActionRedirect) {
+      let transition: TransitionStates["LoadingFetchActionRedirect"] = {
+        location,
+        state,
+        submission: undefined,
+        type: "fetchActionRedirect",
+      };
+      return transition;
     }
   }