From 390c80d7f530cea8473f56dfb6f7d19302d3063f Mon Sep 17 00:00:00 2001
From: Andreas Opferkuch <andreas.opferkuch@gmail.com>
Date: Mon, 6 Jul 2020 12:16:02 +0200
Subject: [PATCH] ThemeContext: Make useStyles type-aware (#26056)

PLUS:
Make it more consise
Add unit test
---
 .../src/themes/ThemeContext.test.tsx          | 27 +++++++++++++++++++
 .../grafana-ui/src/themes/ThemeContext.tsx    |  9 +++----
 2 files changed, 31 insertions(+), 5 deletions(-)
 create mode 100644 packages/grafana-ui/src/themes/ThemeContext.test.tsx

diff --git a/packages/grafana-ui/src/themes/ThemeContext.test.tsx b/packages/grafana-ui/src/themes/ThemeContext.test.tsx
new file mode 100644
index 00000000000..26feb2c617c
--- /dev/null
+++ b/packages/grafana-ui/src/themes/ThemeContext.test.tsx
@@ -0,0 +1,27 @@
+import React from 'react';
+import { config } from '@grafana/runtime';
+import { css } from 'emotion';
+import { mount } from 'enzyme';
+import { useStyles } from './ThemeContext';
+
+describe('useStyles', () => {
+  it('passes in theme and returns style object', () => {
+    const Dummy: React.FC = function() {
+      const styles = useStyles(theme => {
+        expect(theme).toEqual(config.theme);
+
+        return {
+          someStyle: css`
+            color: ${theme?.palette.critical};
+          `,
+        };
+      });
+
+      expect(typeof styles.someStyle).toBe('string');
+
+      return <div>dummy</div>;
+    };
+
+    mount(<Dummy />);
+  });
+});
diff --git a/packages/grafana-ui/src/themes/ThemeContext.tsx b/packages/grafana-ui/src/themes/ThemeContext.tsx
index 0ba6f23e9a0..d5e51752ac9 100644
--- a/packages/grafana-ui/src/themes/ThemeContext.tsx
+++ b/packages/grafana-ui/src/themes/ThemeContext.tsx
@@ -38,12 +38,11 @@ export const withTheme = <P extends Themeable, S extends {} = {}>(Component: Rea
 export function useTheme(): GrafanaTheme {
   return useContext(ThemeContextMock || ThemeContext);
 }
+
 /** Hook for using memoized styles with access to the theme. */
-export const useStyles = (getStyles: (theme?: GrafanaTheme) => any) => {
-  const currentTheme = useTheme();
-  const callback = stylesFactory(stylesTheme => getStyles(stylesTheme));
-  return callback(currentTheme);
-};
+export function useStyles<T>(getStyles: (theme: GrafanaTheme) => T) {
+  return stylesFactory(getStyles)(useTheme());
+}
 
 /**
  * Enables theme context  mocking