From a3c9ea4b24f6dda46c7b7e05acc6dec3f21cb8d4 Mon Sep 17 00:00:00 2001
From: Taylor Lodge <taylor.lodge@koordinates.com>
Date: Thu, 15 Aug 2024 10:08:07 +1200
Subject: [PATCH] fix(testing): allow nesting routing roots inside
 TestRoutingContext

---
 src/routing/index.ts      |  1 +
 src/routing/providers.tsx | 15 ++++++++++++++-
 src/xstateTree.spec.tsx   | 33 +++++++++++++++++++++++++++++++++
 src/xstateTree.tsx        | 10 ++++++++--
 4 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/src/routing/index.ts b/src/routing/index.ts
index 1b5f638..b04fda2 100644
--- a/src/routing/index.ts
+++ b/src/routing/index.ts
@@ -28,5 +28,6 @@ export {
   RoutingContext,
   TestRoutingContext,
   useInRoutingContext,
+  useInTestRoutingContext,
   useActiveRouteEvents,
 } from "./providers";
diff --git a/src/routing/providers.tsx b/src/routing/providers.tsx
index 92666f4..5e96dfe 100644
--- a/src/routing/providers.tsx
+++ b/src/routing/providers.tsx
@@ -5,6 +5,7 @@ import { RoutingEvent } from "./routingEvent";
 
 type Context = {
   activeRouteEvents?: MutableRefObject<RoutingEvent<any>[]>;
+  isTestRoutingContext?: boolean;
 };
 
 export const RoutingContext = createContext<Context | undefined>(undefined);
@@ -30,6 +31,15 @@ export function useInRoutingContext(): boolean {
   return context !== undefined;
 }
 
+/**
+ * @private
+ */
+export function useInTestRoutingContext(): boolean {
+  const context = useContext(RoutingContext);
+
+  return context?.isTestRoutingContext ?? false;
+}
+
 /**
  * @public
  *
@@ -62,7 +72,10 @@ export function TestRoutingContext({
 }) {
   return (
     <RoutingContext.Provider
-      value={{ activeRouteEvents: { current: activeRouteEvents } }}
+      value={{
+        activeRouteEvents: { current: activeRouteEvents },
+        isTestRoutingContext: true,
+      }}
     >
       {children}
     </RoutingContext.Provider>
diff --git a/src/xstateTree.spec.tsx b/src/xstateTree.spec.tsx
index 68dbb99..391533c 100644
--- a/src/xstateTree.spec.tsx
+++ b/src/xstateTree.spec.tsx
@@ -11,6 +11,7 @@ import {
   buildActions,
   createXStateTreeMachine,
 } from "./builders";
+import { TestRoutingContext } from "./routing";
 import { singleSlot } from "./slots";
 import { delay } from "./utils";
 import {
@@ -337,6 +338,38 @@ describe("xstate-tree", () => {
       throw new Error("Should have thrown");
     });
 
+    it("does not throw an error if it's inside a test routing context", async () => {
+      const machine = createMachine({
+        id: "test",
+        initial: "idle",
+        states: {
+          idle: {},
+        },
+      });
+
+      const RootMachine = createXStateTreeMachine(machine, {
+        View() {
+          return <p>I am root</p>;
+        },
+      });
+      const Root = buildRootComponent(RootMachine, {
+        basePath: "/",
+        history: createMemoryHistory(),
+        routes: [],
+      });
+
+      const { rerender } = render(
+        <TestRoutingContext activeRouteEvents={[]}>
+          <Root />
+        </TestRoutingContext>
+      );
+      rerender(
+        <TestRoutingContext activeRouteEvents={[]}>
+          <Root />
+        </TestRoutingContext>
+      );
+    });
+
     it("does not throw an error if either or one are a routing root", async () => {
       const machine = createMachine({
         id: "test",
diff --git a/src/xstateTree.tsx b/src/xstateTree.tsx
index d684318..77cb830 100644
--- a/src/xstateTree.tsx
+++ b/src/xstateTree.tsx
@@ -24,8 +24,9 @@ import {
   RoutingEvent,
   SharedMeta,
   useInRoutingContext,
+  useInTestRoutingContext,
+  useActiveRouteEvents,
 } from "./routing";
-import { useActiveRouteEvents } from "./routing/providers";
 import { GetSlotNames, Slot } from "./slots";
 import { GlobalEvents, AnyXstateTreeMachine, XstateTreeHistory } from "./types";
 import { useConstant } from "./useConstant";
@@ -357,7 +358,12 @@ export function buildRootComponent(
       activeRouteEventsRef.current = events;
     };
     const insideRoutingContext = useInRoutingContext();
-    if (insideRoutingContext && typeof routing !== "undefined") {
+    const inTestRoutingContext = useInTestRoutingContext();
+    if (
+      !inTestRoutingContext &&
+      insideRoutingContext &&
+      typeof routing !== "undefined"
+    ) {
       const m =
         "Routing root rendered inside routing context, this implies a bug";
       if (process.env.NODE_ENV !== "production") {