Enable theme context mocking in tests (#20519)

* Enable theme context mocking in tests

* Expose mockThemeContext from grafana/ui

* Add docs

* Update contribute/style-guides/themes.md

Co-Authored-By: Marcus Olsson <olsson.e.marcus@gmail.com>

* Update packages/grafana-ui/src/themes/ThemeContext.tsx

Co-Authored-By: Marcus Olsson <olsson.e.marcus@gmail.com>

* Update contribute/style-guides/themes.md

Co-Authored-By: Marcus Olsson <olsson.e.marcus@gmail.com>

* Docs update

* Update contribute/style-guides/themes.md

Co-Authored-By: Marcus Olsson <olsson.e.marcus@gmail.com>

contribute/style-guides/themes.md

@@ -29,7 +29,7 @@ const Foo: React.FunctionComponent<FooProps> = () => {
 
 ```
 
-#### Using `withTheme` HOC
+#### Using `withTheme` higher-order component (HOC)
 
 With this method your component will be automatically wrapped in `ThemeContext.Consumer` and provided with current theme via `theme` prop. Component used with `withTheme` must implement `Themeable` interface.
 
@@ -43,6 +43,36 @@ const Foo: React.FunctionComponent<FooProps> = () => ...
 export default withTheme(Foo);
 ```
 
+### Test components that use ThemeContext
+
+When implementing snapshot tests for components that use the `withTheme` HOC, the snapshot will contain the entire theme object. Any change to the theme renders the snapshot outdated.
+
+To make your snapshot theme independent, use the `mockThemeContext` helper function:
+
+```tsx
+import { mockThemeContext } from '@grafana/ui';
+import { MyComponent } from './MyComponent';
+
+describe('MyComponent', () => {
+  let restoreThemeContext;
+
+  beforeAll(() => {
+    // Create ThemeContext mock before any snapshot test is executed
+    restoreThemeContext = mockThemeContext({ type: GrafanaThemeType.Dark });
+  });
+
+  afterAll(() => {
+    // Make sure the theme is restored after snapshot tests are performed
+    restoreThemeContext();
+  });
+
+  it('renders correctyl', () => {
+    const wrapper = mount(<MyComponent />)
+    expect(wrapper).toMatchSnapshot();
+  });
+});
+```
+
 ### Using themes in Storybook
 
 All stories are wrapped with `ThemeContext.Provider` using global decorator. To render `Themeable` component that's not wrapped by `withTheme` HOC you either create a new component in your story:

packages/grafana-ui/src/components/ThresholdsEditor/ThresholdsEditor.test.tsx

 import React, { ChangeEvent } from 'react';
 import { mount } from 'enzyme';
+import { GrafanaThemeType } from '@grafana/data';
 import { ThresholdsEditor, Props, thresholdsWithoutKey } from './ThresholdsEditor';
 import { colors } from '../../utils';
+import { mockThemeContext } from '../../themes/ThemeContext';
 
 const setup = (propOverrides?: Partial<Props>) => {
   const props: Props = {
@@ -25,6 +27,15 @@ function getCurrentThresholds(editor: ThresholdsEditor) {
 }
 
 describe('Render', () => {
+  let restoreThemeContext: any;
+  beforeAll(() => {
+    restoreThemeContext = mockThemeContext({ type: GrafanaThemeType.Dark });
+  });
+
+  afterAll(() => {
+    restoreThemeContext();
+  });
+
   it('should render with base threshold', () => {
     const { wrapper } = setup();
     expect(wrapper).toMatchSnapshot();

packages/grafana-ui/src/themes/ThemeContext.tsx

@@ -3,19 +3,29 @@ import hoistNonReactStatics from 'hoist-non-react-statics';
 
 import { getTheme } from './getTheme';
 import { Themeable } from '../types/theme';
-import { GrafanaThemeType } from '@grafana/data';
+import { GrafanaTheme, GrafanaThemeType } from '@grafana/data';
 
 type Omit<T, K> = Pick<T, Exclude<keyof T, K>>;
 type Subtract<T, K> = Omit<T, keyof K>;
 
+/**
+ * Mock used in tests
+ */
+let ThemeContextMock: React.Context<GrafanaTheme> | null = null;
+
 // Use Grafana Dark theme by default
 export const ThemeContext = React.createContext(getTheme(GrafanaThemeType.Dark));
 ThemeContext.displayName = 'ThemeContext';
 
 export const withTheme = <P extends Themeable, S extends {} = {}>(Component: React.ComponentType<P>) => {
   const WithTheme: React.FunctionComponent<Subtract<P, Themeable>> = props => {
+    /**
+     * If theme context is mocked, let's use it instead of the original context
+     * This is used in tests when mocking theme using mockThemeContext function defined below
+     */
+    const ContextComponent = ThemeContextMock || ThemeContext;
     // @ts-ignore
-    return <ThemeContext.Consumer>{theme => <Component {...props} theme={theme} />}</ThemeContext.Consumer>;
+    return <ContextComponent.Consumer>{theme => <Component {...props} theme={theme} />}</ContextComponent.Consumer>;
   };
 
   WithTheme.displayName = `WithTheme(${Component.displayName})`;
@@ -25,5 +35,15 @@ export const withTheme = <P extends Themeable, S extends {} = {}>(Component: Rea
 };
 
 export function useTheme() {
-  return useContext(ThemeContext);
+  return useContext(ThemeContextMock || ThemeContext);
 }
+
+/**
+ * Enables theme context  mocking
+ */
+export const mockThemeContext = (theme: Partial<GrafanaTheme>) => {
+  ThemeContextMock = React.createContext(theme as GrafanaTheme);
+  return () => {
+    ThemeContextMock = null;
+  };
+};

packages/grafana-ui/src/themes/index.ts

-import { ThemeContext, withTheme, useTheme } from './ThemeContext';
+import { ThemeContext, withTheme, useTheme, mockThemeContext } from './ThemeContext';
 import { getTheme, mockTheme } from './getTheme';
 import { selectThemeVariant } from './selectThemeVariant';
 export { stylesFactory } from './stylesFactory';
-export { ThemeContext, withTheme, mockTheme, getTheme, selectThemeVariant, useTheme };
+export { ThemeContext, withTheme, mockTheme, getTheme, selectThemeVariant, useTheme, mockThemeContext };