Bugfix: Render context around fallback UI so useErrorBoundary works inside of fallback

Author: bvaughn

src/ErrorBoundary.ts

@@ -6,10 +6,7 @@ import {
   PropsWithChildren,
   PropsWithRef,
 } from "react";
-import {
-  ErrorBoundaryContextType,
-  ErrorBoundaryContext,
-} from "./ErrorBoundaryContext";
+import { ErrorBoundaryContext } from "./ErrorBoundaryContext";
 import { ErrorBoundaryProps, FallbackProps } from "./types";
 
 type ErrorBoundaryState = { didCatch: boolean; error: any };
@@ -78,18 +75,20 @@ export class ErrorBoundary extends Component<
       this.props;
     const { didCatch, error } = this.state;
 
+    let childToRender = children;
+
     if (didCatch) {
       const props: FallbackProps = {
         error,
         resetErrorBoundary: this.resetErrorBoundary,
       };
 
       if (isValidElement(fallback)) {
-        return fallback;
+        childToRender = fallback;
       } else if (typeof fallbackRender === "function") {
-        return fallbackRender(props);
+        childToRender = fallbackRender(props);
       } else if (FallbackComponent) {
-        return createElement(FallbackComponent, props);
+        childToRender = createElement(FallbackComponent, props);
       } else {
         throw new Error(
           "react-error-boundary requires either a fallback, fallbackRender, or FallbackComponent prop"
@@ -106,7 +105,7 @@ export class ErrorBoundary extends Component<
           resetErrorBoundary: this.resetErrorBoundary,
         },
       },
-      children
+      childToRender
     );
   }
 }

src/ErrorBoundaryContext.ts

@@ -1,7 +1,7 @@
 import { createContext } from "react";
 
 export type ErrorBoundaryContextType = {
-  didCatch: false;
+  didCatch: boolean;
   error: any;
   resetErrorBoundary: (...args: any[]) => void;
 };

src/useErrorBoundary.test.tsx

@@ -62,4 +62,40 @@ describe("useErrorBoundary", () => {
     });
     expect(container.textContent).toBe("Child");
   });
+
+  it("should work within a fallback component", () => {
+    let resetBoundary: UseErrorBoundaryApi<Error>["resetBoundary"] | null =
+      null;
+    let showBoundary: UseErrorBoundaryApi<Error>["showBoundary"] | null = null;
+
+    function FallbackComponent() {
+      resetBoundary = useErrorBoundary<Error>().resetBoundary;
+      return <div>Error</div>;
+    }
+
+    function Child() {
+      showBoundary = useErrorBoundary<Error>().showBoundary;
+      return <div>Child</div>;
+    }
+
+    const root = createRoot(container);
+    act(() => {
+      root.render(
+        <ErrorBoundary FallbackComponent={FallbackComponent}>
+          <Child />
+        </ErrorBoundary>
+      );
+    });
+    expect(container.textContent).toBe("Child");
+
+    act(() => {
+      showBoundary!(new Error("Example"));
+    });
+    expect(container.textContent).toBe("Error");
+
+    act(() => {
+      resetBoundary!();
+    });
+    expect(container.textContent).toBe("Child");
+  });
 });