From 488c5b8396dc5132236980053867dfcfbf88002e Mon Sep 17 00:00:00 2001
From: Tom Ratcliffe <tom.ratcliffe@grafana.com>
Date: Tue, 9 Apr 2024 13:14:58 +0100
Subject: [PATCH] Use `PropsWithChildren` and allow router configuration in
 test utils render

---
 public/test/test-utils.tsx | 44 +++++++++++++++++---------------------
 1 file changed, 20 insertions(+), 24 deletions(-)

diff --git a/public/test/test-utils.tsx b/public/test/test-utils.tsx
index 2fa8c63daa5..999b7489684 100644
--- a/public/test/test-utils.tsx
+++ b/public/test/test-utils.tsx
@@ -2,7 +2,7 @@ import { ToolkitStore } from '@reduxjs/toolkit/dist/configureStore';
 import { render, RenderOptions } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { KBarProvider } from 'kbar';
-import React from 'react';
+import React, { ComponentProps, Fragment, PropsWithChildren } from 'react';
 import { Provider } from 'react-redux';
 import { MemoryRouter } from 'react-router-dom';
 import { PreloadedState } from 'redux';
@@ -17,25 +17,27 @@ import { configureStore } from 'app/store/configureStore';
 import { StoreState } from 'app/types/store';
 
 interface ExtendedRenderOptions extends RenderOptions {
+  /**
+   * Optional store to use for rendering. If not provided, a fresh store will be generated
+   * via `configureStore` method
+   */
+  store?: ToolkitStore;
   /**
    * Partial state to use for preloading store when rendering tests
    */
   preloadedState?: PreloadedState<StoreState>;
   /**
-   * Optional
+   * Should the wrapper be generated with a wrapping Router component?
+   * Useful if you're testing something that needs more nuanced routing behaviour
+   * and you want full control over it instead
    */
-  store?: ToolkitStore;
   renderWithRouter?: boolean;
+  /**
+   * Props to pass to `MemoryRouter`, if being used
+   */
+  routerOptions?: ComponentProps<typeof MemoryRouter>;
 }
 
-/**
- * Wraps children in empty fragment - used to conditionally render in a router or not,
- * as needed by tests
- */
-const FragmentWrapper = ({ children }: { children: React.ReactNode }) => {
-  return children;
-};
-
 /**
  * Get a wrapper component that implements all of the providers that components
  * within the app will need
@@ -43,22 +45,16 @@ const FragmentWrapper = ({ children }: { children: React.ReactNode }) => {
 const getWrapper = ({
   store,
   renderWithRouter,
+  routerOptions,
   grafanaContext,
-}: {
-  store?: ToolkitStore;
-  /**
-   * Should the wrapper be generated with a wrapping Router component?
-   * Useful if you're testing something that needs more nuanced routing behaviour
-   * and you want full control over it instead
-   */
-  renderWithRouter?: boolean;
+}: ExtendedRenderOptions & {
   grafanaContext?: GrafanaContextType;
-}): React.FC<{ children: React.ReactNode }> => {
+}) => {
   const reduxStore = store || configureStore();
   /**
    * Conditional router - either a MemoryRouter or just a Fragment
    */
-  const PotentialRouter = renderWithRouter ? MemoryRouter : FragmentWrapper;
+  const PotentialRouter = renderWithRouter ? MemoryRouter : Fragment;
 
   const context = {
     ...getGrafanaContextMock(),
@@ -69,14 +65,14 @@ const getWrapper = ({
    * Returns a wrapper that should (closely) match the main `AppWrapper`, so any tests are rendering
    * in mostly the same providers as a "real" hierarchy
    */
-  return function Wrapper({ children }: { children?: React.ReactNode }) {
+  return function Wrapper({ children }: PropsWithChildren) {
     return (
       <Provider store={reduxStore}>
         <ErrorBoundaryAlert style="page">
           <GrafanaContext.Provider value={context}>
             <ThemeProvider value={config.theme2}>
               <KBarProvider>
-                <PotentialRouter>
+                <PotentialRouter {...routerOptions}>
                   <ModalsContextProvider>{children}</ModalsContextProvider>
                 </PotentialRouter>
               </KBarProvider>
@@ -99,7 +95,7 @@ const customRender = (
 ) => {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   const store = renderOptions.preloadedState ? configureStore(renderOptions?.preloadedState as any) : undefined;
-  const AllTheProviders = renderOptions.wrapper || getWrapper({ store, renderWithRouter });
+  const AllTheProviders = renderOptions.wrapper || getWrapper({ store, renderWithRouter, ...renderOptions });
 
   return {
     ...render(ui, { wrapper: AllTheProviders, ...renderOptions }),