diff --git a/.github/workflows/ci_tests.yml b/.github/workflows/ci_tests.yml
index af788d4..e5f6566 100644
--- a/.github/workflows/ci_tests.yml
+++ b/.github/workflows/ci_tests.yml
@@ -50,12 +50,6 @@ jobs:
           name: bqplot-image-gl-dist-${{ github.run_number }}
           path: ./dist
 
-      - name: Install node
-        uses: actions/setup-node@v2
-        with:
-          node-version: '18.x'
-          registry-url: 'https://registry.npmjs.org'
-
       - name: Install Python
         uses: actions/setup-python@v2
         with:
@@ -66,7 +60,7 @@ jobs:
           echo $PWD
           ls -al .
           ls -al dist/
-          pip install dist/bqplot_image_gl*.whl "jupyterlab<4" "pytest-ipywidgets[solara]"
+          pip install dist/bqplot_image_gl*.whl "pytest-ipywidgets[all]"
 
       - name: Install playwright
         run: |
@@ -76,36 +70,12 @@ jobs:
         run: |
           pytest tests/ui
 
-      - name: Install Galata
-        run: |
-          yarn install
-          yarn playwright install chromium
-        working-directory: ui-tests
-
-      - name: Launch JupyterLab
-        run: yarn run start:detached
-        working-directory: ui-tests
-
-      - name: Wait for JupyterLab
-        uses: ifaxity/wait-on-action@v1
-        with:
-          resource: http-get://localhost:8988/api
-          timeout: 20000
-
-      - name: Run UI Tests
-        env:
-          TARGET_URL: http://127.0.0.1:8988
-        run: yarn run test
-        working-directory: ui-tests
-
       - name: Upload UI Test artifacts
         if: always()
         uses: actions/upload-artifact@v3
         with:
-          name: ui-test-output
+          name: bqplot-image-gl-ui-test-output
           path: |
-            ui-tests/playwright-report
-            ui-tests/test-results
             test-results
 
   tests:
diff --git a/tests/ui/image_test.py b/tests/ui/image_test.py
index 71b99e3..fbd79d0 100644
--- a/tests/ui/image_test.py
+++ b/tests/ui/image_test.py
@@ -1,32 +1,35 @@
 import ipywidgets as widgets
 import playwright.sync_api
 from IPython.display import display
-import numpy as np
-from bqplot import Figure, LinearScale, Axis, ColorScale
-from bqplot_image_gl import ImageGL
 
 
-def test_widget_image(solara_test, page_session: playwright.sync_api.Page, assert_solara_snapshot):
 
-    scale_x = LinearScale(min=0, max=1)
-    scale_y = LinearScale(min=0, max=1)
-    scales = {"x": scale_x, "y": scale_y}
-    axis_x = Axis(scale=scale_x, label="x")
-    axis_y = Axis(scale=scale_y, label="y", orientation="vertical")
+def test_widget_image(ipywidgets_runner, page_session: playwright.sync_api.Page, assert_solara_snapshot):
 
-    figure = Figure(scales=scales, axes=[axis_x, axis_y])
+    def kernel_code():
+        import numpy as np
+        from bqplot import Figure, LinearScale, Axis, ColorScale
+        from bqplot_image_gl import ImageGL
+        scale_x = LinearScale(min=0, max=1)
+        scale_y = LinearScale(min=0, max=1)
+        scales = {"x": scale_x, "y": scale_y}
+        axis_x = Axis(scale=scale_x, label="x")
+        axis_y = Axis(scale=scale_y, label="y", orientation="vertical")
 
-    scales_image = {"x": scale_x, "y": scale_y, "image": ColorScale(min=0, max=2)}
+        figure = Figure(scales=scales, axes=[axis_x, axis_y])
 
-    data = np.array([[0., 1.], [2., 3.]])
-    image = ImageGL(image=data, scales=scales_image)
+        scales_image = {"x": scale_x, "y": scale_y, "image": ColorScale(min=0, max=2)}
 
-    figure.marks = (image,)
+        data = np.array([[0., 1.], [2., 3.]])
+        image = ImageGL(image=data, scales=scales_image)
 
-    display(figure)
+        figure.marks = (image,)
 
+        display(figure)
 
+    ipywidgets_runner(kernel_code)
     svg = page_session.locator(".bqplot")
     svg.wait_for()
-    # page_session.wait_for_timeout(1000)
+    # make sure the image is rendered
+    page_session.wait_for_timeout(100)
     assert_solara_snapshot(svg.screenshot())
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_lab-chromium-darwin-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_lab-chromium-darwin-reference.png
new file mode 100644
index 0000000..9f2ea9d
Binary files /dev/null and b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_lab-chromium-darwin-reference.png differ
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_lab-chromium-linux-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_lab-chromium-linux-reference.png
new file mode 100644
index 0000000..3616fb8
Binary files /dev/null and b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_lab-chromium-linux-reference.png differ
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_notebook-chromium-darwin-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_notebook-chromium-darwin-reference.png
new file mode 100644
index 0000000..68d8aa6
Binary files /dev/null and b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_notebook-chromium-darwin-reference.png differ
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_notebook-chromium-linux-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_notebook-chromium-linux-reference.png
new file mode 100644
index 0000000..0d69b69
Binary files /dev/null and b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-jupyter_notebook-chromium-linux-reference.png differ
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-chromium-darwin-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-solara-chromium-darwin-reference.png
similarity index 100%
rename from tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-chromium-darwin-reference.png
rename to tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-solara-chromium-darwin-reference.png
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-chromium-linux-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-solara-chromium-linux-reference.png
similarity index 100%
rename from tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-chromium-linux-reference.png
rename to tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-solara-chromium-linux-reference.png
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-voila-chromium-darwin-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-voila-chromium-darwin-reference.png
new file mode 100644
index 0000000..c6d5855
Binary files /dev/null and b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-voila-chromium-darwin-reference.png differ
diff --git a/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-voila-chromium-linux-reference.png b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-voila-chromium-linux-reference.png
new file mode 100644
index 0000000..e84a424
Binary files /dev/null and b/tests/ui/snapshots/tests/ui/image_test.py/test_widget_image-voila-chromium-linux-reference.png differ
diff --git a/ui-tests/README.md b/ui-tests/README.md
deleted file mode 100644
index 2c40e62..0000000
--- a/ui-tests/README.md
+++ /dev/null
@@ -1,41 +0,0 @@
-# Visual regression tests using Galata
-
-This directory contains visual regression tests for bqplot-image-gl, using Galata.
-
-In order to run them, you need to install dependencies:
-
-```bash
-$ conda install -c conda-forge "yarn<2" jupyterlab=3.5.3
-$ yarn install$
-$ npx playwright install chromium
-```
-
-Then start JupyterLab in one terminal (you need to check that it properly starts on port 8988):
-```bash
-$ yarn run start-jlab
-```
-
-Finally, run the galata tests:
-```bash
-TARGET_URL=http://127.0.0.1:8988 yarn run test
-```
-
-If bqplot-image-gl visuals change, you can re-generate reference images by running:
-```bash
-yarn test:update
-```
-
-## Notebooks directory
-
-The `tests/notebooks` directory contains the test notebooks. For most notebooks (*e.g.* `bars.ipynb`, `scatter.ipynb`) Galata will run them cell by cell and take a screenshot of each output, comparing with the reference images.
-
-When running notebooks named `*_update.ipynb`, Galata will always take the first cell output as reference which must contain the plot, later cells will only be used to update the plot, those notebooks are checking that bqplot-image-gl is properly taking updates into account on already-created plots.
-
-## Add a new test
-
-You can add a new test by simply adding a new notebook to the `tests/notebooks` directory and updating the references. If you want to test updating plots, create notebook named `*_update.ipynb`, create a plot in your first cell then update the plot in later cells.
-
-
-## Updating reference images
-
-In CI, just say 'update galata' (without quotes) in a message to trigger the update of the reference images.
\ No newline at end of file
diff --git a/ui-tests/jupyter_server_config.py b/ui-tests/jupyter_server_config.py
deleted file mode 100644
index 0fba6a7..0000000
--- a/ui-tests/jupyter_server_config.py
+++ /dev/null
@@ -1,11 +0,0 @@
-from tempfile import mkdtemp
-
-c.ServerApp.port = 8988
-c.ServerApp.token = ""
-c.ServerApp.port_retries = 0
-c.ServerApp.password = ""
-c.ServerApp.disable_check_xsrf = True
-c.ServerApp.open_browser = False
-c.ServerApp.root_dir = mkdtemp(prefix='galata-test-')
-
-c.LabApp.expose_app_in_browser = True
diff --git a/ui-tests/package.json b/ui-tests/package.json
deleted file mode 100644
index 6385809..0000000
--- a/ui-tests/package.json
+++ /dev/null
@@ -1,21 +0,0 @@
-{
-  "name": "bqplot-image-gl-ui-tests",
-  "version": "1.0.0",
-  "description": "bqplot-image-gl UI Tests",
-  "private": true,
-  "scripts": {
-    "start": "jupyter lab --config ./jupyter_server_config.py",
-    "start:detached": "yarn run start&",
-    "test": "playwright test",
-    "test:debug": "PWDEBUG=1 playwright test",
-    "test:report": "http-server ./playwright-report -a localhost -o",
-    "test:update": "playwright test --update-snapshots"
-  },
-  "author": "bqplot-image-gl",
-  "license": "Apache-2.0",
-  "dependencies": {
-    "@jupyterlab/galata": "~4.5.0",
-    "klaw-sync": "^6.0.0",
-    "rimraf": "^3.0.2"
-  }
-}
diff --git a/ui-tests/playwright.config.js b/ui-tests/playwright.config.js
deleted file mode 100644
index 6cf25a6..0000000
--- a/ui-tests/playwright.config.js
+++ /dev/null
@@ -1,7 +0,0 @@
-const baseConfig = require('@jupyterlab/galata/lib/playwright-config');
-
-module.exports = {
-  ...baseConfig,
-  timeout: 600000,
-  retries: 1,
-};
diff --git a/ui-tests/tests/bqplot-image-gl.test.ts b/ui-tests/tests/bqplot-image-gl.test.ts
deleted file mode 100644
index fb5774b..0000000
--- a/ui-tests/tests/bqplot-image-gl.test.ts
+++ /dev/null
@@ -1,134 +0,0 @@
-// Copyright (c) Jupyter Development Team.
-// Distributed under the terms of the Modified BSD License.
-
-import { IJupyterLabPageFixture, test } from '@jupyterlab/galata';
-import { expect } from '@playwright/test';
-import * as path from 'path';
-const klaw = require('klaw-sync');
-
-
-const filterUpdateNotebooks = item => {
-  const basename = path.basename(item.path);
-  return basename.includes('_update');
-}
-
-const testCellOutputs = async (page: IJupyterLabPageFixture, tmpPath: string, theme: 'JupyterLab Light' | 'JupyterLab Dark') => {
-  const paths = klaw(path.resolve(__dirname, './notebooks'), {filter: item => !filterUpdateNotebooks(item), nodir: true});
-  const notebooks = paths.map(item => path.basename(item.path));
-
-  const contextPrefix = theme == 'JupyterLab Light' ? 'light' : 'dark';
-  page.theme.setTheme(theme);
-
-  for (const notebook of notebooks) {
-    let results = [];
-
-    await page.notebook.openByPath(`${tmpPath}/${notebook}`);
-    await page.notebook.activate(notebook);
-
-    let numCellImages = 0;
-
-    const getCaptureImageName = (contextPrefix: string, notebook: string, id: number): string => {
-      return `${contextPrefix}-${notebook}-cell-${id}.png`;
-    };
-
-    await page.notebook.runCellByCell({
-      onAfterCellRun: async (cellIndex: number) => {
-        const cell = await page.notebook.getCellOutput(cellIndex);
-        if (cell) {
-          results.push(await cell.screenshot());
-          numCellImages++;
-        }
-      }
-    });
-
-    await page.notebook.save();
-
-    for (let c = 0; c < numCellImages; ++c) {
-      expect(results[c]).toMatchSnapshot(getCaptureImageName(contextPrefix, notebook, c));
-    }
-
-    await page.notebook.close(true);
-  }
-}
-
-const testPlotUpdates = async (page: IJupyterLabPageFixture, tmpPath: string, theme: 'JupyterLab Light' | 'JupyterLab Dark') => {
-  const paths = klaw(path.resolve(__dirname, './notebooks'), {filter: item => filterUpdateNotebooks(item), nodir: true});
-  const notebooks = paths.map(item => path.basename(item.path));
-
-  const contextPrefix = theme == 'JupyterLab Light' ? 'light' : 'dark';
-  page.theme.setTheme(theme);
-
-  for (const notebook of notebooks) {
-    let results = [];
-
-    await page.notebook.openByPath(`${tmpPath}/${notebook}`);
-    await page.notebook.activate(notebook);
-
-    const getCaptureImageName = (contextPrefix: string, notebook: string, id: number): string => {
-      return `${contextPrefix}-${notebook}-cell-${id}.png`;
-    };
-
-    let cellCount = 0;
-    await page.notebook.runCellByCell({
-      onAfterCellRun: async (cellIndex: number) => {
-        // Always get first cell output which must contain the plot
-        const cell = await page.notebook.getCellOutput(0);
-        if (cell) {
-          results.push(await cell.screenshot());
-          cellCount++;
-        }
-      }
-    });
-
-    await page.notebook.save();
-
-    for (let i = 0; i < cellCount; i++) {
-      expect(results[i]).toMatchSnapshot(getCaptureImageName(contextPrefix, notebook, i));
-    }
-
-    await page.notebook.close(true);
-  }
-};
-
-test.describe('bqplot Visual Regression', () => {
-  test.beforeEach(async ({ page, tmpPath }) => {
-    page.on("console", (message) => {
-      console.log('CONSOLE MSG ---', message.text());
-    });
-
-    await page.contents.uploadDirectory(
-      path.resolve(__dirname, './notebooks'),
-      tmpPath
-    );
-    await page.filebrowser.openDirectory(tmpPath);
-  });
-
-  test('Light theme: Check bqplot-image-gl first renders', async ({
-    page,
-    tmpPath,
-  }) => {
-    await testCellOutputs(page, tmpPath, 'JupyterLab Light');
-  });
-
-  // For now we do not test with the dark theme
-  // test('Dark theme: Check bqplot-image-gl first renders', async ({
-  //   page,
-  //   tmpPath,
-  // }) => {
-  //   await testCellOutputs(page, tmpPath, 'JupyterLab Dark');
-  // });
-
-  test('Light theme: Check bqplot-image-gl update plot properties', async ({
-    page,
-    tmpPath,
-  }) => {
-    await testPlotUpdates(page, tmpPath, 'JupyterLab Light');
-  });
-
-  // test('Dark theme: Check bqplot-image-gl update plot properties', async ({
-  //   page,
-  //   tmpPath,
-  // }) => {
-  //   await testPlotUpdates(page, tmpPath, 'JupyterLab Dark');
-  // });
-});
diff --git a/ui-tests/tests/bqplot-image-gl.test.ts-snapshots/light-image-ipynb-cell-0-linux.png b/ui-tests/tests/bqplot-image-gl.test.ts-snapshots/light-image-ipynb-cell-0-linux.png
deleted file mode 100644
index 1365285..0000000
Binary files a/ui-tests/tests/bqplot-image-gl.test.ts-snapshots/light-image-ipynb-cell-0-linux.png and /dev/null differ
diff --git a/ui-tests/tests/notebooks/image.ipynb b/ui-tests/tests/notebooks/image.ipynb
deleted file mode 100644
index ae894bd..0000000
--- a/ui-tests/tests/notebooks/image.ipynb
+++ /dev/null
@@ -1,56 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "df77670d",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import numpy as np\n",
-    "from bqplot import Figure, LinearScale, Axis, ColorScale\n",
-    "from bqplot_image_gl import ImageGL\n",
-    "scale_x = LinearScale(min=0, max=1, allow_padding=False)\n",
-    "scale_y = LinearScale(min=0, max=1, allow_padding=False)\n",
-    "scales = {'x': scale_x,\n",
-    "          'y': scale_y}\n",
-    "axis_x = Axis(scale=scale_x, label='x', )\n",
-    "axis_y = Axis(scale=scale_y, label='y', orientation='vertical')\n",
-    "\n",
-    "figure = Figure(scales=scales, axes=[axis_x, axis_y])\n",
-    "\n",
-    "scales_image = {'x': scale_x,\n",
-    "                'y': scale_y,\n",
-    "                'image': ColorScale(min=0, max=1)}\n",
-    "\n",
-    "s = 1\n",
-    "np.random.seed(0)\n",
-    "image = ImageGL(image=np.random.random((10, 10)).astype('float32'), scales=scales_image)\n",
-    "\n",
-    "figure.marks = (image,)\n",
-    "figure"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.9.16"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}