diff --git a/package.json b/package.json
index fb59177..37b4a80 100644
--- a/package.json
+++ b/package.json
@@ -11,8 +11,10 @@
     "push": "clasp push",
     "setup:https": "mkdirp certs && mkcert -key-file ./certs/key.pem -cert-file ./certs/cert.pem localhost 127.0.0.1",
     "build:dev": "tsc && vite build --mode development",
+    "build:test": "tsc && NODE_ENV=test vite build --mode test",
     "build": "tsc && vite build --mode production",
     "deploy:dev": "yarn build:dev && yarn push",
+    "deploy:test": "yarn build:test && yarn push",
     "deploy": "yarn build && yarn push",
     "start": "yarn deploy:dev && yarn dev"
   },
diff --git a/src/client/create-diagram-dialog/components/create-diagram-dialog.tsx b/src/client/create-diagram-dialog/components/create-diagram-dialog.tsx
index 166662d..8722f5e 100644
--- a/src/client/create-diagram-dialog/components/create-diagram-dialog.tsx
+++ b/src/client/create-diagram-dialog/components/create-diagram-dialog.tsx
@@ -10,7 +10,6 @@ const CreateDiagramDialog = () => {
 
   useEffect(() => {
     if (!authState?.authorized) return;
-    // const url = buildUrl('/app/plugins/confluence/select', state.token);
     const url = buildUrl(
       '/app/diagrams/new?pluginSource=googledocs',
       authState.token
diff --git a/src/client/create-diagram-dialog/index.jsx b/src/client/create-diagram-dialog/index.jsx
index c85839d..31b2420 100644
--- a/src/client/create-diagram-dialog/index.jsx
+++ b/src/client/create-diagram-dialog/index.jsx
@@ -1,7 +1,7 @@
-import ReactDOM from 'react-dom';
+import { createRoot } from 'react-dom/client';
 import CreateDiagramDialog from './components/create-diagram-dialog';
 import './styles.css';
 
 const container = document.getElementById('index');
-const root = ReactDOM.createRoot(container);
+const root = createRoot(container);
 root.render(<CreateDiagramDialog />);
diff --git a/src/client/edit-diagram-dialog/index.jsx b/src/client/edit-diagram-dialog/index.jsx
index 0102351..d1f605d 100644
--- a/src/client/edit-diagram-dialog/index.jsx
+++ b/src/client/edit-diagram-dialog/index.jsx
@@ -1,8 +1,7 @@
-import ReactDOM from 'react-dom';
+import { createRoot } from 'react-dom/client';
 import EditDiagramDialog from './components/edit-diagram-dialog';
-
 import './styles.css';
 
 const container = document.getElementById('index');
-const root = ReactDOM.createRoot(container);
+const root = createRoot(container);
 root.render(<EditDiagramDialog />);
diff --git a/src/client/hooks/useAuth.ts b/src/client/hooks/useAuth.ts
index 51e7711..492fa6a 100644
--- a/src/client/hooks/useAuth.ts
+++ b/src/client/hooks/useAuth.ts
@@ -1,5 +1,6 @@
 import { useCallback, useEffect, useState } from 'react';
 import { serverFunctions } from '../utils/serverFunctions';
+import { baseURL } from '../../config/urls';
 
 type Status = 'idle' | 'loading' | 'success' | 'error';
 
@@ -21,6 +22,7 @@ const useAuth = () => {
   const getAuth = useCallback(async () => {
     setAuthStatus('loading');
     try {
+      await serverFunctions.setBaseUrl(baseURL);
       const state = await serverFunctions.getAuthorizationState();
       setAuthState(state as AuthState);
       setAuthStatus('success');
diff --git a/src/client/preview-diagram-dialog/index.jsx b/src/client/preview-diagram-dialog/index.jsx
index 44b3dc1..9dcd08e 100644
--- a/src/client/preview-diagram-dialog/index.jsx
+++ b/src/client/preview-diagram-dialog/index.jsx
@@ -1,7 +1,7 @@
-import ReactDOM from 'react-dom';
+import { createRoot } from 'react-dom/client';
 import PreviewDiagramDialog from './components/preview-diagram-dialog';
 import './styles.css';
 
 const container = document.getElementById('index');
-const root = ReactDOM.createRoot(container);
+const root = createRoot(container);
 root.render(<PreviewDiagramDialog />);
diff --git a/src/client/select-diagram-dialog/components/select-diagram-dialog.tsx b/src/client/select-diagram-dialog/components/select-diagram-dialog.tsx
index 799fb08..d30bfb6 100644
--- a/src/client/select-diagram-dialog/components/select-diagram-dialog.tsx
+++ b/src/client/select-diagram-dialog/components/select-diagram-dialog.tsx
@@ -10,7 +10,6 @@ const SelectDiagramDialog = () => {
 
   useEffect(() => {
     if (!authState?.authorized) return;
-    // const url = buildUrl('/app/plugins/confluence/select', state.token);
     const url = buildUrl(
       '/app/plugins/select?pluginSource=googledocs',
       authState.token
diff --git a/src/client/select-diagram-dialog/index.jsx b/src/client/select-diagram-dialog/index.jsx
index 0a0d4f3..caa570d 100644
--- a/src/client/select-diagram-dialog/index.jsx
+++ b/src/client/select-diagram-dialog/index.jsx
@@ -1,7 +1,7 @@
-import ReactDOM from 'react-dom';
+import { createRoot } from 'react-dom/client';
 import SelectDiagramDialog from './components/select-diagram-dialog';
 import './styles.css';
 
 const container = document.getElementById('index');
-const root = ReactDOM.createRoot(container);
+const root = createRoot(container);
 root.render(<SelectDiagramDialog />);
diff --git a/src/client/sidebar/components/Sidebar.tsx b/src/client/sidebar/components/Sidebar.tsx
index 69e9e13..2d20165 100644
--- a/src/client/sidebar/components/Sidebar.tsx
+++ b/src/client/sidebar/components/Sidebar.tsx
@@ -33,7 +33,6 @@ const Sidebar = () => {
 
   useEffect(() => {
     if (!authState?.authorized) return;
-    // const url = buildUrl('/app/plugins/confluence/select', state.token);
     const url = buildUrl(
       '/app/plugins/recent?pluginSource=googledocs',
       authState.token
@@ -65,9 +64,11 @@ const Sidebar = () => {
   useEffect(() => {
     const handleMessage = async (e: MessageEvent) => {
       const action = e.data.action;
-      console.log('action', action);
+      const actionData = e.data;
+
       if (action === 'save') {
-        const data = e.data.data;
+        const data = actionData.data;
+        if (!data) return;
         const metadata = new URLSearchParams({
           projectID: data.projectID,
           documentID: data.documentID,
@@ -83,20 +84,24 @@ const Sidebar = () => {
         } catch (error) {
           console.error('Error inserting image with metadata', error);
         }
-      } else if (action === 'edit') {
-        const data = e.data;
-        if (!data.editUrl) return;
+        return;
+      }
+      if (action === 'edit') {
+        const editUrl = actionData.editUrl;
+        if (!editUrl) return;
         try {
-          localStorage.setItem('editUrl', data.editUrl);
+          localStorage.setItem('editUrl', editUrl);
           await serverFunctions.openEditDiagramDialogWithUrl();
         } catch (error) {
           console.error('Error opening edit dialog', error);
         }
-      } else if (action === 'view') {
-        console.log(e.data);
-        if (!e.data.url) return;
+        return;
+      }
+      if (action === 'view') {
+        const viewUrl = actionData.url;
+        if (!viewUrl) return;
         try {
-          localStorage.setItem('previewUrl', e.data.url);
+          localStorage.setItem('previewUrl', viewUrl);
           await serverFunctions.openPreviewDiagramDialog();
         } catch (error) {
           console.error('Error opening edit dialog', error);
diff --git a/src/client/utils/helpers.ts b/src/client/utils/helpers.ts
index 2696144..a2025ff 100644
--- a/src/client/utils/helpers.ts
+++ b/src/client/utils/helpers.ts
@@ -1,19 +1,9 @@
-import { baseURL } from '../../utils/urls';
-
-interface Document {
-  documentID: string;
-  major: string;
-  minor: string;
-}
+import { baseURL } from '../../config/urls';
 
 export const buildUrl = (pathname: string, accessToken: string) => {
   return `${baseURL}/oauth/frame?token=${accessToken}&redirect=${pathname}`;
 };
 
-export const buildRawUrl = (document: Document, theme = 'light') => {
-  return `${baseURL}/raw/${document.documentID}?version=v${document.major}.${document.minor}&theme=${theme}&format=png`;
-};
-
 export const handleDialogClose = () => {
   if ((window as any).google) {
     (window as any).google.script.host.close();
diff --git a/src/config/urls.ts b/src/config/urls.ts
new file mode 100644
index 0000000..1c03ed9
--- /dev/null
+++ b/src/config/urls.ts
@@ -0,0 +1,3 @@
+export const devUrl = 'https://test.mermaidchart.com';
+export const prodUrl = 'https://mermaidchart.com';
+export const baseURL = (import.meta as any).env.PROD ? prodUrl : devUrl;
diff --git a/src/server/index.ts b/src/server/index.ts
index 8050ee9..15afe64 100644
--- a/src/server/index.ts
+++ b/src/server/index.ts
@@ -7,6 +7,7 @@ import {
   openPreviewDiagramDialog,
   openSidebar,
   getOAuthURL,
+  setBaseUrl,
   handleCallback,
   getAuthorizationState,
   resetOAuth,
@@ -28,6 +29,7 @@ export {
   openPreviewDiagramDialog,
   openSidebar,
   getOAuthURL,
+  setBaseUrl,
   handleCallback,
   getAuthorizationState,
   resetOAuth,
diff --git a/src/server/ui.js b/src/server/ui.js
index 4dbfb4b..c293aab 100644
--- a/src/server/ui.js
+++ b/src/server/ui.js
@@ -1,5 +1,3 @@
-const baseURL = 'https://test.mermaidchart.com';
-
 export const onOpen = () => {
   const menu = DocumentApp.getUi()
     .createAddonMenu()
@@ -11,6 +9,26 @@ export const onOpen = () => {
   menu.addToUi();
 };
 
+export const setBaseUrl = (url) => {
+  try {
+    const scriptProperties = PropertiesService.getScriptProperties();
+    scriptProperties.setProperty('baseURL', url);
+  } catch (err) {
+    Logger.log('Failed with error %s', err.message);
+  }
+};
+
+const getBaseUrl = () => {
+  try {
+    const scriptProperties = PropertiesService.getScriptProperties();
+    const baseURL = scriptProperties.getProperty('baseURL');
+    Logger.log('Base URL: %s', baseURL);
+    return baseURL;
+  } catch (error) {
+    Logger.log('Failed with error %s', error.message);
+  }
+};
+
 export const openCreateDiagramDialog = () => {
   const html = HtmlService.createHtmlOutputFromFile('create-diagram-dialog')
     .append(
@@ -125,6 +143,11 @@ export function handleCallback(callbackRequest) {
 function getOAuthService() {
   pkceChallengeVerifier();
   const userProps = PropertiesService.getUserProperties();
+  const baseURL = getBaseUrl();
+
+  if (!baseURL) {
+    throw new Error('Base URL is not defined.');
+  }
 
   return OAuth2.createService('Mermaid Chart')
     .setAuthorizationBaseUrl(baseURL + '/oauth/authorize')
@@ -349,6 +372,11 @@ export function syncImages(maxWidth = 400) {
   const { token } = getAuthorizationState();
   const body = DocumentApp.getActiveDocument().getBody();
   const images = body.getImages();
+  const baseURL = getBaseUrl();
+
+  if (!baseURL) {
+    throw new Error('Base URL is not defined.');
+  }
 
   images.forEach((image) => {
     const altDescription = image.getAltDescription();
diff --git a/src/utils/urls.js b/src/utils/urls.js
deleted file mode 100644
index c1957f9..0000000
--- a/src/utils/urls.js
+++ /dev/null
@@ -1 +0,0 @@
-export const baseURL = 'https://test.mermaidchart.com';
diff --git a/vite.config.ts b/vite.config.ts
index 3231126..c3e096b 100644
--- a/vite.config.ts
+++ b/vite.config.ts
@@ -164,16 +164,16 @@ const buildConfig = ({ mode }: { mode: string }) => {
         targets,
       }),
       /**
-       * This builds the client react app bundles for production, and writes them to disk.
+       * This builds the client react app bundles for production and test, and writes them to disk.
        * Because multiple client entrypoints (dialogs) are built, we need to loop through
        * each entrypoint and build the client bundle for each. Vite doesn't have great tooling for
        * building multiple single-page apps in one project, so we have to do this manually with a
        * post-build closeBundle hook (https://rollupjs.org/guide/en/#closebundle).
        */
-      mode === 'production' && {
-        name: 'build-client-production-bundles',
+      (mode === 'production' || mode === 'test') && {
+        name: `build-client-${mode}-bundles`,
         closeBundle: async () => {
-          console.log('Building client production bundles...');
+          console.log(`Building client ${mode} bundles...`);
           // eslint-disable-next-line no-restricted-syntax
           for (const clientEntrypoint of clientEntrypoints) {
             console.log('Building client bundle for', clientEntrypoint.name);