From bf3d58fd08dc1c44a3468eb76e4c46da9e3c69f6 Mon Sep 17 00:00:00 2001
From: Brijesh Bittu <brijeshb42@gmail.com>
Date: Mon, 3 Jun 2024 09:37:53 +0530
Subject: [PATCH] Review changes

---
 .../src/utils/generateThemeSource.ts          | 22 +++------
 packages/pigment-css-react/tests/testUtils.ts |  2 +-
 .../useTheme/generateThemeSource.test.ts      | 45 +++++++++++++++++++
 .../tests/useTheme/useTheme.test.tsx          |  2 +-
 4 files changed, 52 insertions(+), 19 deletions(-)
 create mode 100644 packages/pigment-css-react/tests/useTheme/generateThemeSource.test.ts

diff --git a/packages/pigment-css-react/src/utils/generateThemeSource.ts b/packages/pigment-css-react/src/utils/generateThemeSource.ts
index 2db3b6e3..b2abcaa4 100644
--- a/packages/pigment-css-react/src/utils/generateThemeSource.ts
+++ b/packages/pigment-css-react/src/utils/generateThemeSource.ts
@@ -1,24 +1,12 @@
 import type { Theme } from './extendTheme';
 import { generateThemeTokens } from './generateCss';
 
-export function generateThemeSource(theme: Theme | undefined) {
+export function generateThemeSource(theme?: Theme) {
   if (!theme) {
     return `export default {}`;
   }
-  const hasBreakpoints = !!theme.unstable_args?.breakpoints;
-  const hasTransitions = !!theme.unstable_args?.transitions;
-  const breakpointsImport = hasBreakpoints
-    ? "import createBreakpoints from '@mui/system/createTheme/createBreakpoints';\n"
-    : '';
-  const transitionsImport = hasTransitions
-    ? `import createTransitions from '${process.env.PACKAGE_NAME}/createTransitions';\n`
-    : '';
-  const source = `${breakpointsImport}${transitionsImport}
-const theme = ${JSON.stringify(generateThemeTokens(theme))};
-${hasBreakpoints ? `theme.breakpoints = createBreakpoints(${JSON.stringify(theme.unstable_args.breakpoints)});\n` : ''}
-${hasTransitions ? `theme.transitions = createTransitions(${JSON.stringify(theme.unstable_args.transitions)});\n` : ''}
-export default theme;
-`;
-
-  return source;
+  if (typeof theme.toRuntimeSource !== 'function') {
+    return `export default ${JSON.stringify(generateThemeTokens(theme))};`;
+  }
+  return theme.toRuntimeSource.call(theme, theme);
 }
diff --git a/packages/pigment-css-react/tests/testUtils.ts b/packages/pigment-css-react/tests/testUtils.ts
index ae1b7208..846e602c 100644
--- a/packages/pigment-css-react/tests/testUtils.ts
+++ b/packages/pigment-css-react/tests/testUtils.ts
@@ -56,7 +56,7 @@ export async function runTransformation(
       if (source !== '@pigment-css/react' && !source.endsWith('/zero-styled')) {
         return null;
       }
-      return require.resolve(`${pkgJson['wyw-in-js'].tags[tag]}`);
+      return require.resolve(`../${pkgJson['wyw-in-js'].tags[tag]}`.replace('.js', ''));
     },
   };
 
diff --git a/packages/pigment-css-react/tests/useTheme/generateThemeSource.test.ts b/packages/pigment-css-react/tests/useTheme/generateThemeSource.test.ts
new file mode 100644
index 00000000..a9802809
--- /dev/null
+++ b/packages/pigment-css-react/tests/useTheme/generateThemeSource.test.ts
@@ -0,0 +1,45 @@
+import { expect } from 'chai';
+import { generateThemeSource } from '../../utils';
+
+describe('Pigment CSS - generateThemeSource', () => {
+  it('should export empty theme if theme is undefined', () => {
+    expect(generateThemeSource()).to.equal(`export default {}`);
+  });
+
+  it('should generate basic theme code with theme json', () => {
+    const result = generateThemeSource({
+      vars: {
+        palette: {
+          primary: {
+            main: 'red',
+            secondary: 'blue',
+          },
+        },
+      },
+    });
+
+    expect(result).to.equal(
+      `export default {"vars":{"palette":{"primary":{"main":"red","secondary":"blue"}}}};`,
+    );
+  });
+
+  it('should generate theme from toRuntimeSource method if present', () => {
+    const result = generateThemeSource({
+      vars: {
+        palette: {
+          primary: {
+            main: 'red',
+            secondary: 'blue',
+          },
+        },
+      },
+      toRuntimeSource(theme: any) {
+        return `const theme = ${JSON.stringify(theme)};export default theme;`;
+      },
+    });
+
+    expect(result).to.equal(
+      `const theme = {"vars":{"palette":{"primary":{"main":"red","secondary":"blue"}}}};export default theme;`,
+    );
+  });
+});
diff --git a/packages/pigment-css-react/tests/useTheme/useTheme.test.tsx b/packages/pigment-css-react/tests/useTheme/useTheme.test.tsx
index 1aa20267..6011a673 100644
--- a/packages/pigment-css-react/tests/useTheme/useTheme.test.tsx
+++ b/packages/pigment-css-react/tests/useTheme/useTheme.test.tsx
@@ -2,7 +2,7 @@ import path from 'node:path';
 import { runTransformation, expect } from '../testUtils';
 
 describe('Pigment CSS - useTheme', () => {
-  it('basics', async () => {
+  it('should replace useTheme call with the theme import', async () => {
     const { output, fixture } = await runTransformation(
       path.join(__dirname, 'fixtures/useTheme.input.js'),
     );