From ffe3450739027bcf5f38e566a4e484c7cb7d57e2 Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Fri, 3 Jul 2020 14:43:14 +1000
Subject: [PATCH] Properly mock location in app stories

---
 lib/router/src/router.tsx  |  2 ++
 lib/ui/src/app.stories.tsx | 27 +++++++++++++--------------
 lib/ui/src/index.tsx       |  7 ++++---
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/lib/router/src/router.tsx b/lib/router/src/router.tsx
index 669b4cfc8fb..9092898555f 100644
--- a/lib/router/src/router.tsx
+++ b/lib/router/src/router.tsx
@@ -9,6 +9,7 @@ import {
   RouteComponentProps,
   LocationContext,
   NavigateFn,
+  History,
 } from '@reach/router';
 import { ToggleVisibility } from './visibility';
 import { queryFromString, parsePath, getMatch, StoryData } from './utils';
@@ -106,3 +107,4 @@ export { QueryLocation as Location };
 export { Route };
 export { queryNavigate as navigate };
 export { LocationProvider };
+export { History };
diff --git a/lib/ui/src/app.stories.tsx b/lib/ui/src/app.stories.tsx
index 126f36868e4..6f575d4a137 100644
--- a/lib/ui/src/app.stories.tsx
+++ b/lib/ui/src/app.stories.tsx
@@ -13,19 +13,18 @@ export default {
   parameters: {
     layout: 'fullscreen',
   },
-  decorators: [
-    ((StoryFn, c) => (
-      <LocationProvider
-        key="location.provider"
-        history={createHistory(createMemorySource('/?path=/story/story--id'))}
-      >
-        <StoryFn {...c} />
-      </LocationProvider>
-    )) as DecoratorFn,
-  ],
 };
 
-export const Default = () => <App provider={(new FakeProvider() as unknown) as Provider} />;
-export const LoadingState = () => (
-  <App provider={(new PrettyFakeProvider() as unknown) as Provider} />
-);
+const Story = (args: React.ComponentProps<typeof App>) => <App {...args} />;
+
+export const Default = Story.bind({});
+Default.args = {
+  provider: (new FakeProvider() as unknown) as Provider,
+  history: createHistory(createMemorySource('/?path=/story/story--id')),
+};
+
+export const LoadingState = Story.bind({});
+LoadingState.args = {
+  ...Default.args,
+  provider: (new PrettyFakeProvider() as unknown) as Provider,
+};
diff --git a/lib/ui/src/index.tsx b/lib/ui/src/index.tsx
index 082896e77cf..8bce4381eec 100644
--- a/lib/ui/src/index.tsx
+++ b/lib/ui/src/index.tsx
@@ -2,7 +2,7 @@ import { DOCS_MODE } from 'global';
 import React, { FunctionComponent } from 'react';
 import ReactDOM from 'react-dom';
 
-import { Location, LocationProvider } from '@storybook/router';
+import { Location, LocationProvider, History } from '@storybook/router';
 import { Provider as ManagerProvider, Combo } from '@storybook/api';
 import { ThemeProvider, ensure as ensureTheme } from '@storybook/theming';
 import { HelmetProvider } from 'react-helmet-async';
@@ -28,12 +28,13 @@ const Container = process.env.XSTORYBOOK_EXAMPLE_APP ? React.StrictMode : React.
 
 export interface RootProps {
   provider: Provider;
+  history: History;
 }
 
-export const Root: FunctionComponent<RootProps> = ({ provider }) => (
+export const Root: FunctionComponent<RootProps> = ({ provider, history }) => (
   <Container key="container">
     <HelmetProvider key="helmet.Provider">
-      <LocationProvider key="location.provider">
+      <LocationProvider key="location.provider" history={history}>
         <Location key="location.consumer">
           {(locationData) => (
             <ManagerProvider