diff --git a/.github/workflows/storybook-tests.yml b/.github/workflows/storybook-tests.yml
index bc79392..f309b9b 100644
--- a/.github/workflows/storybook-tests.yml
+++ b/.github/workflows/storybook-tests.yml
@@ -35,16 +35,22 @@ jobs:
       - name: Build Storybook
         run: yarn build-storybook
       
+      - name: Install test dependencies
+        run: yarn add --dev http-server wait-on
+
       - name: Run Storybook tests
         run: |
           # Start Storybook in the background
-          npx http-server storybook-static --port 6006 --silent &
+          yarn http-server storybook-static --port 6006 --silent &
           
           # Wait for Storybook to be ready
-          npx wait-on http://localhost:6006
+          yarn wait-on http://localhost:6006
+          
+          # Set environment variable for Storybook URL
+          export STORYBOOK_URL=http://localhost:6006
           
-          # Run the tests using test-runner directly
-          npx test-storybook --url http://localhost:6006 --browsers chromium
+          # Run the tests using Vitest
+          yarn test:storybook
         
       - name: Upload test artifacts on failure
         if: failure()
diff --git a/.storybook/main.ts b/.storybook/main.ts
index 9e12745..21e318d 100644
--- a/.storybook/main.ts
+++ b/.storybook/main.ts
@@ -5,6 +5,7 @@ const config: StorybookConfig = {
   addons: [
     '@storybook/addon-essentials',
     '@storybook/addon-interactions',
+    '@storybook/addon-vitest'
   ],
   framework: {
     name: '@storybook/html-vite',
diff --git a/.storybook/vitest.setup.ts b/.storybook/vitest.setup.ts
new file mode 100644
index 0000000..c583eaf
--- /dev/null
+++ b/.storybook/vitest.setup.ts
@@ -0,0 +1,6 @@
+import { setProjectAnnotations } from '@storybook/html-vite';
+import * as projectAnnotations from './preview';
+
+// This is an important step to apply the right configuration when testing your stories.
+// More info at: https://storybook.js.org/docs/api/portable-stories/portable-stories-vitest#setprojectannotations
+setProjectAnnotations([projectAnnotations]);
\ No newline at end of file
diff --git a/.windsurf/rules/storybook.md b/.windsurf/rules/storybook.md
index 9f41380..961de9c 100644
--- a/.windsurf/rules/storybook.md
+++ b/.windsurf/rules/storybook.md
@@ -6,4 +6,6 @@ This is guidance for generation and maintenance of StoryBook stories.
 
 Lots of data tables are generated for StoryBook tests.  Do not create them as template text in the story.  Generate them as suitably named `.html` files in the `stories` sub-folder.
 
-This make it easier to understand and maintain them.
\ No newline at end of file
+This make it easier to understand and maintain them.
+
+We will be using Storybook V9 interaction tests, as documented here: https://storybook.js.org/docs/writing-tests/interaction-testing
\ No newline at end of file
diff --git a/package.json b/package.json
index 82a9695..9940757 100644
--- a/package.json
+++ b/package.json
@@ -7,28 +7,31 @@
     "dev": "vite",
     "build": "tsc && vite build && node scripts/copy-demo.js",
     "preview": "vite preview",
-    "preview:demo": "vite preview --port 3000 --open /demo",
+    "preview:demo": "vite preview --port 3000 --open grid-sight",
     "storybook": "storybook dev -p 6006",
     "build-storybook": "storybook build",
-    "test:storybook": "test-storybook",
+    "test:storybook": "vitest run --project=storybook",
     "test": "vitest run",
     "test:watch": "vitest",
     "test:coverage": "vitest run --coverage",
-    "test:e2e": "playwright test tests/e2e/demo.spec.ts",
-    "test:e2e:ui": "playwright test --ui tests/e2e/demo.spec.ts"
+    "test:e2e": "vite build && playwright test tests/e2e",
+    "test:e2e:ui": "vite build && playwright test --ui tests/e2e"
   },
   "devDependencies": {
     "@playwright/test": "^1.53.0",
     "@storybook/addon-docs": "^9.0.6",
     "@storybook/addon-essentials": "^9.0.0-alpha.12",
     "@storybook/addon-interactions": "^9.0.0-alpha.10",
+    "@storybook/addon-vitest": "^9.0.9",
     "@storybook/html-vite": "^9.0.6",
     "@storybook/test": "^8.6.14",
     "@types/fs-extra": "^11.0.4",
+    "@vitest/browser": "3.2.3",
     "@vitest/coverage-v8": "^3.2.3",
     "@vitest/ui": "^3.2.3",
     "fs-extra": "^11.3.0",
     "jsdom": "^26.1.0",
+    "playwright": "^1.53.0",
     "storybook": "^9.0.6",
     "terser": "^5.42.0",
     "typescript": "~5.8.3",
diff --git a/public/demo/index.html b/public/index.html
similarity index 75%
rename from public/demo/index.html
rename to public/index.html
index 5ac46f0..299eb74 100644
--- a/public/demo/index.html
+++ b/public/index.html
@@ -60,58 +60,34 @@ <h2>Basic Usage</h2>
       <div class="table-wrapper">
         <h3>Product Inventory</h3>
         <table class="inventory-table">
-          <thead>
-            <tr>
-              <th>Product ID</th>
-              <th>Product Name</th>
-              <th>Category</th>
-              <th>Price</th>
-              <th>In Stock</th>
-              <th>Last Ordered</th>
-            </tr>
-          </thead>
-          <tbody>
-            <tr>
-              <td>P1001</td>
-              <td>Wireless Mouse</td>
-              <td>Electronics</td>
-              <td>29.99</td>
-              <td>45</td>
-              <td>2023-05-15</td>
-            </tr>
-            <tr>
-              <td>P1002</td>
-              <td>Ergonomic Keyboard</td>
-              <td>Electronics</td>
-              <td>89.99</td>
-              <td>23</td>
-              <td>2023-05-18</td>
-            </tr>
-            <tr>
-              <td>P2001</td>
-              <td>Desk Lamp</td>
-              <td>Furniture</td>
-              <td>45.50</td>
-              <td>12</td>
-              <td>2023-05-10</td>
-            </tr>
-            <tr>
-              <td>P2002</td>
-              <td>Office Chair</td>
-              <td>Furniture</td>
-              <td>249.99</td>
-              <td>5</td>
-              <td>2023-05-05</td>
-            </tr>
-            <tr>
-              <td>P3001</td>
-              <td>Notebook Set</td>
-              <td>Stationery</td>
-              <td>12.99</td>
-              <td>87</td>
-              <td>2023-05-20</td>
-            </tr>
-          </tbody>
+          <tr>
+            <td>Gear / Revs</td>
+            <td>5</td>
+            <td>10</td>
+            <td>15</td>
+            <td>20</td>
+          </tr>
+          <tr>
+            <td>Low</td>
+            <td>10</td>
+            <td>15</td>
+            <td>20</td>
+            <td>25</td>
+          </tr>
+          <tr>
+            <td>Medium</td>
+            <td>17</td>
+            <td>22</td>
+            <td>27</td>
+            <td>32</td>
+          </tr>
+          <tr>
+            <td>High</td>
+            <td>27</td>
+            <td>25</td>
+            <td>30</td>
+            <td>35</td>
+          </tr>
         </table>
       </div>
 
@@ -124,7 +100,7 @@ <h3>Monthly Sales Data</h3>
               <th>Revenue</th>
               <th>Expenses</th>
               <th>Profit</th>
-              <th>Growth %</th>
+              <th>Dominant item</th>
             </tr>
           </thead>
           <tbody>
@@ -133,35 +109,35 @@ <h3>Monthly Sales Data</h3>
               <td>12500</td>
               <td>9800</td>
               <td>2700</td>
-              <td>5.2</td>
+              <td>Corded Mouse</td>
             </tr>
             <tr>
               <td>February</td>
               <td>11800</td>
               <td>9200</td>
               <td>2600</td>
-              <td>-3.7</td>
+              <td>Wireless Mouse</td>
             </tr>
             <tr>
               <td>March</td>
               <td>14200</td>
               <td>10100</td>
               <td>4100</td>
-              <td>20.3</td>
+              <td>Corded Mouse</td>
             </tr>
             <tr>
               <td>April</td>
               <td>13800</td>
               <td>9900</td>
               <td>3900</td>
-              <td>-2.8</td>
+              <td>Corded Mouse</td>
             </tr>
             <tr>
               <td>May</td>
               <td>15600</td>
               <td>10500</td>
               <td>5100</td>
-              <td>13.0</td>
+              <td>Desk Lamp</td>
             </tr>
           </tbody>
         </table>
@@ -182,7 +158,7 @@ <h3>Browser Console</h3>
   </footer>
 
   <!-- Load Grid-Sight library -->
-  <script src="grid-sight.iife.js"></script>
+  <script src="../dist/grid-sight.iife.js"></script>
   
   <script>
     // Simple console logger for the demo
diff --git a/scripts/copy-demo.js b/scripts/copy-demo.js
index 4945427..91fae89 100644
--- a/scripts/copy-demo.js
+++ b/scripts/copy-demo.js
@@ -12,7 +12,7 @@ const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
 
 const rootDir = path.resolve(__dirname, '..');
-const sourceDir = path.join(rootDir, 'public', 'demo');
+const sourceDir = path.join(rootDir, 'public');
 const targetDir = path.join(rootDir, 'dist');
 
 async function copyDemoFiles() {
diff --git a/src/__tests__/enrichments/frequency.test.ts b/src/__tests__/enrichments/frequency.test-interactions.ts
similarity index 100%
rename from src/__tests__/enrichments/frequency.test.ts
rename to src/__tests__/enrichments/frequency.test-interactions.ts
diff --git a/src/__tests__/utils/frequency.test.ts b/src/__tests__/utils/frequency-calcs.test.ts
similarity index 100%
rename from src/__tests__/utils/frequency.test.ts
rename to src/__tests__/utils/frequency-calcs.test.ts
diff --git a/src/enrichments/__tests__/heatmap.test.ts b/src/enrichments/__tests__/heatmap-interactions.test.ts
similarity index 89%
rename from src/enrichments/__tests__/heatmap.test.ts
rename to src/enrichments/__tests__/heatmap-interactions.test.ts
index 2e94296..8c65e54 100644
--- a/src/enrichments/__tests__/heatmap.test.ts
+++ b/src/enrichments/__tests__/heatmap-interactions.test.ts
@@ -6,28 +6,34 @@ describe('Heatmap', () => {
   const createTestTable = (): HTMLTableElement => {
     const table = document.createElement('table');
     table.innerHTML = `
-      <thead>
-        <tr>
-          <th>Product</th>
-          <th>Q1</th>
-          <th>Q2</th>
-          <th>Q3</th>
-        </tr>
-      </thead>
-      <tbody>
-        <tr>
-          <td>Widget A</td>
-          <td>10</td>
-          <td>20</td>
-          <td>30</td>
-        </tr>
-        <tr>
-          <td>Widget B</td>
-          <td>15</td>
-          <td>25</td>
-          <td>35</td>
-        </tr>
-      </tbody>
+  <tr>
+    <td>Gear / Revs</td>
+    <td>5</td>
+    <td>10</td>
+    <td>15</td>
+    <td>20</td>
+  </tr>
+  <tr>
+    <td>Low</td>
+    <td id='cell-1-1'>10</td>
+    <td>15</td>
+    <td>20</td>
+    <td>25</td>
+  </tr>
+  <tr>
+    <td>Medium</td>
+    <td>17</td>
+    <td>22</td>
+    <td>27</td>
+    <td>32</td>
+  </tr>
+  <tr>
+    <td>High</td>
+    <td>27</td>
+    <td>25</td>
+    <td>30</td>
+    <td>35</td>
+  </tr>
     `;
     document.body.appendChild(table);
     return table;
@@ -124,7 +130,7 @@ describe('Heatmap', () => {
       toggleHeatmap(table, columnIndex, 'column');
       
       // Get a reference to a cell that should be colored
-      const cell = table.querySelector('td') as HTMLElement;
+      const cell = table.querySelector('#cell-1-1') as HTMLElement;
       expect(cell).toBeTruthy();
       const initialColor = cell.style.backgroundColor;
       expect(initialColor).toBeTruthy();
@@ -144,11 +150,11 @@ describe('Heatmap', () => {
       toggleHeatmap(table, 1, 'column');
       
       // Get the cell that will be split
-      const cell = table.querySelector('tbody tr:nth-child(1) td:nth-child(1)') as HTMLElement;
+      const cell = table.querySelector('#cell-1-1') as HTMLElement;
       const initialColor = cell.style.backgroundColor;
-      
+
       // Apply row heatmap that will overlap with the column
-      toggleHeatmap(table, 0, 'row');
+      toggleHeatmap(table, 2, 'row');
       
       // Should have the split class and custom properties
       expect(cell.classList.contains('gs-heatmap-split')).toBe(true);
@@ -159,8 +165,8 @@ describe('Heatmap', () => {
       expect(cell.style.backgroundColor).toBe('');
       
       // Remove the row heatmap
-      removeHeatmap(table, 0, 'row');
-      
+      removeHeatmap(table, 2, 'row');
+     
       // Should remove the split and restore column heatmap
       expect(cell.classList.contains('gs-heatmap-split')).toBe(false);
       expect(cell.style.backgroundColor).toBe(initialColor);
@@ -175,14 +181,14 @@ describe('Heatmap', () => {
 
       // Apply both row and column heatmaps
       toggleHeatmap(table, 1, 'column');
-      toggleHeatmap(table, 0, 'row');
+      toggleHeatmap(table, 2, 'row');
       
       // Get the cell at the intersection
-      const cell = table.querySelector('tbody tr:nth-child(1) td:nth-child(2)') as HTMLElement;
-      
+      const cell = table.querySelector('#cell-1-1') as HTMLElement;
+
       // Remove the column heatmap
       removeHeatmap(table, 1, 'column');
-      
+
       // Should remove the split and restore row heatmap
       expect(cell.classList.contains('gs-heatmap-split')).toBe(false);
       expect(cell.style.getPropertyValue('--split-color-1')).toBe('');
@@ -198,7 +204,7 @@ describe('Heatmap', () => {
     it('should apply heatmap to a numeric row', () => {
       // Arrange
       const table = createTestTable();
-      const rowIndex = 1; // First data row
+      const rowIndex = 2; // First data row
       
       // Act - apply heatmap to the first data row
       toggleHeatmap(table, rowIndex, 'row');
@@ -208,7 +214,7 @@ describe('Heatmap', () => {
       expect(table.classList.contains('gs-heatmap')).toBe(true);
       
       // 2. Check if cells in the row have background colors
-      const row = table.rows[rowIndex];
+      const row = table.rows[rowIndex - 1];
       expect(row).toBeTruthy();
       
       // Skip the first cell (text label) and check the rest
@@ -245,7 +251,7 @@ describe('Heatmap', () => {
       
       // Apply both row and column heatmaps to create a split cell
       applyHeatmap(table, 1, 'column');
-      applyHeatmap(table, 1, 'row');
+      applyHeatmap(table, 2, 'row');
       
       // Get the cell at the intersection
       const cell = getCell(table, 1, 1);
@@ -262,9 +268,8 @@ describe('Heatmap', () => {
       expect(cell.classList.contains('gs-heatmap-split')).toBe(false);
       expect(cell.style.getPropertyValue('--split-color-1')).toBe('');
       expect(cell.style.getPropertyValue('--split-color-2')).toBe('');
-      
+
       // The other heatmap should still be active
-      expect(cell.classList.contains('gs-heatmap-cell')).toBe(true);
       expect(cell.style.backgroundColor).toBeTruthy();
       
       // Clean up
diff --git a/src/enrichments/heatmap.ts b/src/enrichments/heatmap.ts
index 0574c0b..f884224 100644
--- a/src/enrichments/heatmap.ts
+++ b/src/enrichments/heatmap.ts
@@ -497,20 +497,21 @@ export function toggleHeatmap(
   type: HeatmapType = 'column',
   options: HeatmapOptions = {}
 ): void {
-  if (type === 'table') {
-    // For table-wide heatmap, delegate to applyTableHeatmap
-    applyTableHeatmap(table, options);
-    return;
-  }
-  
   // Check if this specific heatmap is already active
-  const isActive = isHeatmapActive(table, index, type);
+  // For table type, use index -1
+  const effectiveIndex = type === 'table' ? -1 : index
+  const isActive = isHeatmapActive(table, effectiveIndex, type)
   
   if (isActive) {
     // If it's active, remove just this specific heatmap
-    removeHeatmap(table, index, type);
+    removeHeatmap(table, effectiveIndex, type)
   } else {
-    // If not active, apply the new heatmap (can have multiple heatmaps on different rows/columns)
-    applyHeatmap(table, index, type, options);
+    // If not active, apply the new heatmap
+    if (type === 'table') {
+      applyTableHeatmap(table, options)
+    } else {
+      // For row/column, apply the new heatmap
+      applyHeatmap(table, index, type, options)
+    }
   }
 }
diff --git a/src/ui/frequency-chart-dialog.ts b/src/ui/frequency-chart-dialog.ts
index 1974aea..6d24958 100644
--- a/src/ui/frequency-chart-dialog.ts
+++ b/src/ui/frequency-chart-dialog.ts
@@ -1,5 +1,11 @@
 import type { FrequencyResult } from '../utils/frequency';
 
+// Color palette for frequency chart bars
+const BAR_COLORS = [
+  '#4e79a7', '#f28e2c', '#e15759', '#76b7b2', '#59a14f',
+  '#edc949', '#af7aa1', '#ff9da7', '#9c755f', '#bab0ab'
+];
+
 // CSS class names
 const DIALOG_CLASS = 'gs-frequency-chart-dialog';
 const DIALOG_VISIBLE_CLASS = 'gs-frequency-chart-dialog--visible';
@@ -97,12 +103,11 @@ const DIALOG_STYLES = `
 }
 
 .${DIALOG_SVG_BAR_CLASS} {
-  fill: #0066cc;
   transition: fill 0.2s ease;
 }
 
 .${DIALOG_SVG_BAR_CLASS}:hover {
-  fill: #004c99;
+  opacity: 0.8;
 }
 
 .${DIALOG_SVG_LABEL_CLASS} {
@@ -347,6 +352,10 @@ export class FrequencyChartDialog {
       bar.setAttribute('height', barHeight.toString());
       bar.setAttribute('class', DIALOG_SVG_BAR_CLASS);
       
+      // Assign different colors to each bar
+      const colorIndex = index % BAR_COLORS.length;
+      bar.setAttribute('fill', BAR_COLORS[colorIndex]);
+      
       // Add tooltip with data
       bar.setAttribute('data-value', value);
       bar.setAttribute('data-count', count.toString());
diff --git a/src/ui/toggle-injector.ts b/src/ui/toggle-injector.ts
index 1eae54c..7f2a157 100644
--- a/src/ui/toggle-injector.ts
+++ b/src/ui/toggle-injector.ts
@@ -2,7 +2,7 @@ import { injectPlusIcons, removePlusIcons, plusIconStyles } from './header-utils
 import type { HeaderType } from './header-utils';
 import { removeAllMenus } from './enrichment-menu';
 import { analyzeTable } from '../core/table-detection';
-import { toggleHeatmap, applyTableHeatmap } from '../enrichments/heatmap';
+import { toggleHeatmap } from '../enrichments/heatmap';
 import { calculateStatistics } from '../enrichments/statistics';
 import { analyzeFrequencies } from '../utils/frequency';
 import { cleanNumericCell } from '../core/type-detection';
@@ -134,20 +134,57 @@ function handleEnrichmentSelected(event: Event) {
         toggleHeatmap(table, rowIndex + 1, 'row');
       }
     } else if (type === 'table') {
-      // Apply heatmap to all numeric cells in the table
-      applyTableHeatmap(table);
+      // Toggle heatmap on all numeric cells in the table
+      toggleHeatmap(table, -1, 'table');
     }
-  } else if (enrichmentType === 'statistics' && type === 'column') {
-    // Type assertion for table header cell
-    const th = header as HTMLTableCellElement;
-    const columnIndex = th.cellIndex;
-    if (columnIndex >= 0) {
+  } else if (enrichmentType === 'statistics') {
+    if (type === 'column') {
+      // Type assertion for table header cell
+      const th = header as HTMLTableCellElement;
+      const columnIndex = th.cellIndex;
+      if (columnIndex >= 0) {
+        try {
+          const values = extractNumericColumnValues(table, columnIndex);
+          const stats = calculateStatistics(values);
+          window._gsStatisticsPopup.show(stats, header);
+        } catch (error) {
+          console.error('Error calculating statistics:', error);
+          alert(`Error: ${error instanceof Error ? error.message : 'Unknown error'}`);
+        }
+      }
+    } else if (type === 'row') {
+      // Type assertion for table row
+      const tr = header.closest('tr') as HTMLTableRowElement;
+      if (!tr) {
+        console.error('Could not find row');
+        return;
+      }
+      
+      try {
+        const values = extractNumericRowValues(tr);
+        if (values.length === 0) {
+          alert('No numeric values found in this row');
+          return;
+        }
+        
+        const stats = calculateStatistics(values);
+        window._gsStatisticsPopup.show(stats, header);
+      } catch (error) {
+        console.error('Error calculating statistics for row:', error);
+        alert(`Error: ${error instanceof Error ? error.message : 'Unknown error'}`);
+      }
+    } else if (type === 'table') {
       try {
-        const values = extractNumericColumnValues(table, columnIndex);
+        const values = extractNumericTableValues(table);
+        if (values.length === 0) {
+          alert('No numeric values found in this table');
+          return;
+        }
+        
         const stats = calculateStatistics(values);
         window._gsStatisticsPopup.show(stats, header);
       } catch (error) {
-        console.error('Error calculating statistics:', error);
+        console.error('Error calculating statistics for table:', error);
         alert(`Error: ${error instanceof Error ? error.message : 'Unknown error'}`);
       }
     }
@@ -268,6 +305,64 @@ function extractNumericColumnValues(table: HTMLTableElement, columnIndex: number
   return values;
 }
 
+/**
+ * Extracts numeric values from a table row
+ */
+function extractNumericRowValues(row: HTMLTableRowElement): number[] {
+  const values: number[] = [];
+  
+  // Get all cells in the row
+  const cells = Array.from(row.cells);
+  
+  // Skip the first cell if it's a header
+  const startIndex = cells.length > 0 && cells[0].tagName.toLowerCase() === 'th' ? 1 : 0;
+  
+  for (let i = startIndex; i < cells.length; i++) {
+    const value = cleanNumericCell(cells[i].textContent || '');
+    if (value !== null) {
+      values.push(value);
+    }
+  }
+  
+  return values;
+}
+
+/**
+ * Extracts all numeric values from a table, excluding headers
+ */
+function extractNumericTableValues(table: HTMLTableElement): number[] {
+  const values: number[] = [];
+  
+  // Get all rows in the table
+  const rows = Array.from(table.rows);
+  
+  // Skip the header row(s)
+  // If there's a thead, skip all rows in it
+  // Otherwise, skip the first row as it's likely a header
+  const startIndex = table.tHead ? table.tHead.rows.length : 1;
+  
+  // Process all non-header rows
+  for (let i = startIndex; i < rows.length; i++) {
+    const row = rows[i];
+    
+    // Get all cells in the row
+    const cells = Array.from(row.cells);
+    
+    // Skip the first cell if it's a header
+    const cellStartIndex = cells.length > 0 && cells[0].tagName.toLowerCase() === 'th' ? 1 : 0;
+    
+    // Process all non-header cells
+    for (let j = cellStartIndex; j < cells.length; j++) {
+      const value = cleanNumericCell(cells[j].textContent || '');
+      if (value !== null) {
+        values.push(value);
+      }
+    }
+  }
+  
+  return values;
+}
+
 export function injectToggle(table: HTMLTableElement): boolean {
   // Find the first cell in the first row of the thead
   const firstRow = table.tHead?.rows[0] || table.rows[0];
diff --git a/src/utils/frequency.ts b/src/utils/frequency.ts
index 691a5b5..2cb3b19 100644
--- a/src/utils/frequency.ts
+++ b/src/utils/frequency.ts
@@ -59,7 +59,7 @@ function analyzeFrequencies(values: unknown[]): FrequencyResult[] {
   return sortedFrequencies.map(([value, count]) => [
     value,
     count,
-    (count / total) * 100
+    Number(((count / total) * 100).toFixed(1))
   ])
 }
 
diff --git a/tests/e2e/demo.spec.ts b/tests/e2e/demo.spec.ts
index f5af548..6db7406 100644
--- a/tests/e2e/demo.spec.ts
+++ b/tests/e2e/demo.spec.ts
@@ -24,7 +24,7 @@ test.describe('Grid-Sight Demo', () => {
 
     try {
       // Navigate to the demo page
-      await page.goto('http://localhost:3001/demo/');
+      await page.goto('http://localhost:3001/');
       
       // Wait for the page to load completely
       await page.waitForLoadState('domcontentloaded');
@@ -75,7 +75,7 @@ test.describe('Grid-Sight Demo', () => {
     test.skip(!!process.env.CI, 'Skipping file:// test in CI environment');
     
     // Navigate to the demo page directly
-    await page.goto(`file://${process.cwd()}/dist/demo/index.html`);
+    await page.goto(`file://${process.cwd()}/dist/index.html`);
     
     // Wait for the page to load completely
     await page.waitForLoadState('domcontentloaded');
diff --git a/tests/e2e/heatmap.spec.ts b/tests/e2e/heatmap.spec.ts
new file mode 100644
index 0000000..308ed3b
--- /dev/null
+++ b/tests/e2e/heatmap.spec.ts
@@ -0,0 +1,134 @@
+import { test, expect } from '@playwright/test'
+
+/**
+ * End-to-end tests for Grid-Sight heatmap functionality
+ * 
+ * These tests verify that heatmap enrichments can be applied and removed correctly
+ */
+
+test.describe('Grid-Sight Heatmap', () => {
+  // Test setup - start server and navigate to demo page
+  test.beforeEach(async ({ page }) => {
+    // Start the preview server
+    const { preview } = await import('vite')
+    const server = await preview({
+      preview: {
+        port: 3001,
+        open: false,
+      },
+      build: {
+        outDir: 'dist'
+      }
+    })
+
+    // Store the server in the test info
+    test.info().attachments.push({
+      name: 'server',
+      body: Buffer.from(JSON.stringify(server as unknown as Record<string, unknown>)),
+      contentType: 'application/json'
+    })
+
+    // Navigate to the demo page
+    await page.goto('http://localhost:3001/')
+    
+    // Wait for the page to load completely
+    await page.waitForLoadState('domcontentloaded')
+    
+    // Wait for Grid-Sight to initialize
+    await page.waitForFunction(() => {
+      const consoleEl = document.getElementById('console')
+      return consoleEl && consoleEl.textContent && 
+             consoleEl.textContent.includes('GridSight loaded successfully')
+    }, { timeout: 5000 })
+  })
+
+  // Close the server after each test
+  test.afterEach(async ({ page: _page }, testInfo) => {
+    // Get the server from the test info
+    const serverAttachment = testInfo.attachments.find(a => a.name === 'server')
+    if (serverAttachment && serverAttachment.body) {
+      const server = JSON.parse(serverAttachment.body.toString())
+      // Close the preview server - safely handle different server structures
+      if (server.httpServer && typeof server.httpServer.close === 'function') {
+        await new Promise(resolve => server.httpServer.close(resolve))
+      } else if (server.close && typeof server.close === 'function') {
+        await new Promise(resolve => server.close(resolve))
+      }
+    }
+  })
+
+  // Test table-wide heatmap toggle functionality
+  test('should toggle table-wide heatmap on and off correctly', async ({ page }) => {
+    // Get the first table
+    const table = page.locator('table').first()
+    
+    // Find the Grid-Sight toggle button
+    await page.waitForSelector('.grid-sight-toggle')
+    
+    // Click on the Grid-Sight toggle
+    await page.click('.grid-sight-toggle')
+    
+    // await for the table plus to appear
+    await page.waitForSelector('.gs-table-plus', { timeout: 500 })
+
+    // Click on the table plus
+    await page.click('.gs-table-plus')
+
+    // Wait for the menu to appear
+    await page.waitForSelector('.gs-enrichment-menu', { timeout: 500 })
+    
+    // Click on the heatmap option
+    await page.getByText('Heatmap').click()
+    
+    // Wait for the heatmap to be applied
+    await page.waitForFunction(() => {
+      const table = document.querySelector('table')
+      return table && table.classList.contains('gs-heatmap')
+    }, { timeout: 5000 })
+    
+    // Verify that the table has the heatmap class
+    await expect(table).toHaveClass(/gs-heatmap/)
+    
+    // Verify that cells have background colors (indicating heatmap is applied)
+    const cellsWithColor = await page.evaluate(() => {
+      const cells = Array.from(document.querySelectorAll('table td'))
+      return cells.filter(cell => {
+        const style = window.getComputedStyle(cell)
+        return style.backgroundColor !== 'rgba(0, 0, 0, 0)' && 
+               style.backgroundColor !== 'rgb(255, 255, 255)' &&
+               style.backgroundColor !== 'transparent'
+      }).length
+    })
+    
+    // Ensure we have cells with background color
+    expect(cellsWithColor).toBeGreaterThan(0)
+    
+    // Toggle the heatmap off by clicking the menu option again
+    await page.click('.gs-plus-icon')
+    await page.waitForSelector('.gs-enrichment-menu', { timeout: 5000 })
+    await page.getByText('Heatmap').click()
+    
+    // Wait for the heatmap to be removed
+    await page.waitForFunction(() => {
+      const table = document.querySelector('table')
+      return table && !table.classList.contains('gs-heatmap')
+    }, { timeout: 5000 })
+    
+    // Verify that the table no longer has the heatmap class
+    await expect(table).not.toHaveClass(/gs-heatmap/)
+    
+    // Verify that cells no longer have background colors
+    const cellsWithColorAfterToggle = await page.evaluate(() => {
+      const cells = Array.from(document.querySelectorAll('table td'))
+      return cells.filter(cell => {
+        const style = window.getComputedStyle(cell)
+        return style.backgroundColor !== 'rgba(0, 0, 0, 0)' && 
+               style.backgroundColor !== 'rgb(255, 255, 255)' &&
+               style.backgroundColor !== 'transparent'
+      }).length
+    })
+    
+    // Ensure no cells have background color
+    expect(cellsWithColorAfterToggle).toBe(0)
+  })
+})
\ No newline at end of file
diff --git a/vite.config.ts b/vite.config.ts
index 702c711..baa625f 100644
--- a/vite.config.ts
+++ b/vite.config.ts
@@ -5,17 +5,6 @@ import { dirname } from 'path';
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 
-// Simple plugin to copy the demo folder from public to dist
-function copyDemoFolder() {
-  return {
-    name: 'copy-demo-folder',
-    writeBundle() {
-      // This will be implemented in the build script
-      console.log('Demo files will be copied to the dist root in the build script');
-    },
-  };
-}
-
 export default defineConfig({
   // Configure how Vite handles different file types
   assetsInclude: ['**/*.html'],
@@ -23,7 +12,7 @@ export default defineConfig({
   // Configure the development server (optional, for demo serving)
   server: {
     port: 3000,
-    open: '/demo.html',
+    open: '/index.html',
   },
   
   // Configure build options
@@ -68,7 +57,6 @@ export default defineConfig({
   
   // Configure plugins
   plugins: [
-    copyDemoFolder(),
     // Add any other Vite plugins here
   ],
 
diff --git a/vitest.config.ts b/vitest.config.ts
index 3e17b8b..b9a30a8 100644
--- a/vitest.config.ts
+++ b/vitest.config.ts
@@ -6,6 +6,8 @@ export default defineConfig({
     globals: true,
     environment: 'jsdom',
     setupFiles: ['./src/test-setup.ts'],
+    include: ['tests/unit/**/*.test.ts', 'src/**/*.test.ts'],
+    exclude: ['node_modules/**', 'dist/**', '.git/**', 'coverage/**', 'tests/e2e/**'],
     coverage: {
       reporter: ['text', 'json', 'html'],
       exclude: [
diff --git a/vitest.setup.ts b/vitest.setup.ts
new file mode 100644
index 0000000..1690152
--- /dev/null
+++ b/vitest.setup.ts
@@ -0,0 +1,12 @@
+// Setup file for Vitest tests
+import '@testing-library/jest-dom'
+import { vi } from 'vitest'
+
+// Mock browser APIs that might not be available in the test environment
+global.ResizeObserver = vi.fn().mockImplementation(() => ({
+  observe: vi.fn(),
+  unobserve: vi.fn(),
+  disconnect: vi.fn(),
+}))
+
+// Add any other global mocks or setup needed for your tests
diff --git a/vitest.shims.d.ts b/vitest.shims.d.ts
new file mode 100644
index 0000000..f923d47
--- /dev/null
+++ b/vitest.shims.d.ts
@@ -0,0 +1 @@
+/// <reference types="@vitest/browser/providers/playwright" />
\ No newline at end of file
diff --git a/vitest.storybook.config.ts b/vitest.storybook.config.ts
new file mode 100644
index 0000000..b836fcb
--- /dev/null
+++ b/vitest.storybook.config.ts
@@ -0,0 +1,28 @@
+import { defineConfig, mergeConfig } from 'vitest/config'
+import { storybookTest } from '@storybook/addon-vitest/vitest-plugin'
+import viteConfig from './vite.config'
+import path from 'node:path'
+
+export default mergeConfig(
+  viteConfig,
+  defineConfig({
+    plugins: [
+      storybookTest({
+        configDir: path.join(process.cwd(), '.storybook'),
+        storybookScript: 'yarn storybook --ci',
+      }),
+    ],
+    test: {
+      environment: 'jsdom',
+      globals: true,
+      setupFiles: ['./vitest.setup.ts'],
+      include: ['src/**/*.stories.@(js|jsx|ts|tsx)'],
+      browser: {
+        enabled: true,
+        provider: 'playwright',
+        headless: true,
+        instances: [{ browser: 'chromium' }],
+      },
+    },
+  })
+)
diff --git a/vitest.workspace.ts b/vitest.workspace.ts
new file mode 100644
index 0000000..46f2a01
--- /dev/null
+++ b/vitest.workspace.ts
@@ -0,0 +1,32 @@
+import path from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+import { defineWorkspace } from 'vitest/config';
+
+import { storybookTest } from '@storybook/addon-vitest/vitest-plugin';
+
+const dirname =
+  typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+
+// More info at: https://storybook.js.org/docs/next/writing-tests/integrations/vitest-addon
+export default defineWorkspace([
+  'vitest.config.ts',
+  {
+    extends: 'vite.config.ts',
+    plugins: [
+      // The plugin will run tests for the stories defined in your Storybook config
+      // See options at: https://storybook.js.org/docs/next/writing-tests/integrations/vitest-addon#storybooktest
+      storybookTest({ configDir: path.join(dirname, '.storybook') }),
+    ],
+    test: {
+      name: 'storybook',
+      browser: {
+        enabled: true,
+        headless: true,
+        provider: 'playwright',
+        instances: [{ browser: 'chromium' }]
+      },
+      setupFiles: ['.storybook/vitest.setup.ts'],
+    },
+  },
+]);
diff --git a/yarn.lock b/yarn.lock
index 6e843f2..61c6b91 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -469,6 +469,16 @@
     "@storybook/global" "^5.0.0"
     ts-dedent "^2.0.0"
 
+"@storybook/addon-vitest@^9.0.9":
+  version "9.0.9"
+  resolved "https://registry.yarnpkg.com/@storybook/addon-vitest/-/addon-vitest-9.0.9.tgz#06460bf0c69faae174b83301e6d28f7997acc41c"
+  integrity sha512-OVGK1mVlbnlLnZmUNw9guweuTAPg3qsxD9kaWuYNSKytMBx3ELqHbCrFL1lhFdlwXPV3e2t5AK/9PQbEtDYZUA==
+  dependencies:
+    "@storybook/global" "^5.0.0"
+    "@storybook/icons" "^1.4.0"
+    prompts "^2.4.0"
+    ts-dedent "^2.2.0"
+
 "@storybook/builder-vite@9.0.9":
   version "9.0.9"
   resolved "https://registry.yarnpkg.com/@storybook/builder-vite/-/builder-vite-9.0.9.tgz#3073a191718c42181d6fb00c8de2a9823ece9636"
@@ -505,7 +515,7 @@
     "@storybook/global" "^5.0.0"
     ts-dedent "^2.0.0"
 
-"@storybook/icons@^1.2.12":
+"@storybook/icons@^1.2.12", "@storybook/icons@^1.4.0":
   version "1.4.0"
   resolved "https://registry.yarnpkg.com/@storybook/icons/-/icons-1.4.0.tgz#7cf7ab3dfb41943930954c4ef493a73798d8b31d"
   integrity sha512-Td73IeJxOyalzvjQL+JXx72jlIYHgs+REaHiREOqfpo3A2AYYG71AUbcv+lg7mEDIweKVCxsMQ0UKo634c8XeA==
@@ -536,7 +546,7 @@
     "@vitest/expect" "2.0.5"
     "@vitest/spy" "2.0.5"
 
-"@testing-library/dom@10.4.0":
+"@testing-library/dom@10.4.0", "@testing-library/dom@^10.4.0":
   version "10.4.0"
   resolved "https://registry.yarnpkg.com/@testing-library/dom/-/dom-10.4.0.tgz#82a9d9462f11d240ecadbf406607c6ceeeff43a8"
   integrity sha512-pemlzrSESWbdAloYml3bAJMEfNh1Z7EduzqPKprCH5S341frlpYnUEW0H72dLxa6IsYr+mPno20GiSm+h9dEdQ==
@@ -647,6 +657,20 @@
   dependencies:
     csstype "^3.0.2"
 
+"@vitest/browser@3.2.3":
+  version "3.2.3"
+  resolved "https://registry.yarnpkg.com/@vitest/browser/-/browser-3.2.3.tgz#703c1a12e65d6864ba91fc617caa29243c5fa535"
+  integrity sha512-5HpUb0ixGF8JWSAjb/P1x/VPuTYUkL4pL0+YO6DJiuvQgqJN3PREaUEcXwfXjU4nBc37EahfpRbAwdE9pHs9lQ==
+  dependencies:
+    "@testing-library/dom" "^10.4.0"
+    "@testing-library/user-event" "^14.6.1"
+    "@vitest/mocker" "3.2.3"
+    "@vitest/utils" "3.2.3"
+    magic-string "^0.30.17"
+    sirv "^3.0.1"
+    tinyrainbow "^2.0.0"
+    ws "^8.18.2"
+
 "@vitest/coverage-v8@^3.2.3":
   version "3.2.3"
   resolved "https://registry.yarnpkg.com/@vitest/coverage-v8/-/coverage-v8-3.2.3.tgz#866f6d9b7d394dc614f3f2186dbe46b1bc940813"
@@ -1339,6 +1363,11 @@ jsonfile@^6.0.1:
   optionalDependencies:
     graceful-fs "^4.1.6"
 
+kleur@^3.0.3:
+  version "3.0.3"
+  resolved "https://registry.yarnpkg.com/kleur/-/kleur-3.0.3.tgz#a79c9ecc86ee1ce3fa6206d1216c501f147fc07e"
+  integrity sha512-eTIzlVOSUR+JxdDFepEYcBMtZ9Qqdef+rnzWdRZuMbOywu5tO2w2N7rqjoANZ5k9vywhL6Br1VRjUIgTQx4E8w==
+
 lodash@^4.17.21:
   version "4.17.21"
   resolved "https://registry.yarnpkg.com/lodash/-/lodash-4.17.21.tgz#679591c564c3bffaae8454cf0b3df370c3d6911c"
@@ -1490,7 +1519,7 @@ playwright-core@1.53.0:
   resolved "https://registry.yarnpkg.com/playwright-core/-/playwright-core-1.53.0.tgz#bf2738cc143116b6130b78e0c644edf2e7e53ff4"
   integrity sha512-mGLg8m0pm4+mmtB7M89Xw/GSqoNC+twivl8ITteqvAndachozYe2ZA7srU6uleV1vEdAHYqjq+SV8SNxRRFYBw==
 
-playwright@1.53.0:
+playwright@1.53.0, playwright@^1.53.0:
   version "1.53.0"
   resolved "https://registry.yarnpkg.com/playwright/-/playwright-1.53.0.tgz#4eae78aa24e3c714bf71981f80b3310b838692fd"
   integrity sha512-ghGNnIEYZC4E+YtclRn4/p6oYbdPiASELBIYkBXfaTVKreQUYbMUYQDwS12a8F0/HtIjr/CkGjtwABeFPGcS4Q==
@@ -1524,6 +1553,14 @@ pretty-format@^27.0.2:
     ansi-styles "^5.0.0"
     react-is "^17.0.1"
 
+prompts@^2.4.0:
+  version "2.4.2"
+  resolved "https://registry.yarnpkg.com/prompts/-/prompts-2.4.2.tgz#7b57e73b3a48029ad10ebd44f74b01722a4cb069"
+  integrity sha512-NxNv/kLguCA7p3jE8oL2aEBsrJWgAakBpgmgK6lpPWV+WuOmY6r2/zbAVnP+T8bQlA0nzHXSJSJW0Hq7ylaD2Q==
+  dependencies:
+    kleur "^3.0.3"
+    sisteransi "^1.0.5"
+
 punycode@^2.3.1:
   version "2.3.1"
   resolved "https://registry.yarnpkg.com/punycode/-/punycode-2.3.1.tgz#027422e2faec0b25e1549c3e1bd8309b9133b6e5"
@@ -1657,6 +1694,11 @@ sirv@^3.0.1:
     mrmime "^2.0.0"
     totalist "^3.0.0"
 
+sisteransi@^1.0.5:
+  version "1.0.5"
+  resolved "https://registry.yarnpkg.com/sisteransi/-/sisteransi-1.0.5.tgz#134d681297756437cc05ca01370d3a7a571075ed"
+  integrity sha512-bLGGlR1QxBcynn2d5YmDX4MGjlZvy2MRBDRNHLJ8VI6l6+9FUiyTFNJ0IveOSP0bcXgVDPRcfGqA0pjaqUpfVg==
+
 source-map-js@^1.2.0, source-map-js@^1.2.1:
   version "1.2.1"
   resolved "https://registry.yarnpkg.com/source-map-js/-/source-map-js-1.2.1.tgz#1ce5650fddd87abc099eda37dcff024c2667ae46"
@@ -2031,7 +2073,7 @@ wrap-ansi@^8.1.0:
     string-width "^5.0.1"
     strip-ansi "^7.0.1"
 
-ws@^8.18.0:
+ws@^8.18.0, ws@^8.18.2:
   version "8.18.2"
   resolved "https://registry.yarnpkg.com/ws/-/ws-8.18.2.tgz#42738b2be57ced85f46154320aabb51ab003705a"
   integrity sha512-DMricUmwGZUVr++AEAe2uiVM7UoO9MAVZMDu05UQOaUII0lp+zOzLLU4Xqh/JvTqklB1T4uELaaPBKyjE1r4fQ==