ThemeContext: Make useStyles type-aware (#26056)

PLUS:
Make it more consise
Add unit test

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 />);
+  });
+});

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) => {
+export function useStyles<T>(getStyles: (theme: GrafanaTheme) => T) {
-  const currentTheme = useTheme();
+  return stylesFactory(getStyles)(useTheme());
-  const callback = stylesFactory(stylesTheme => getStyles(stylesTheme));
+}
-  return callback(currentTheme);
-};
 
 /**
  * Enables theme context  mocking