diff --git a/.dockerignore b/.dockerignore
index 684c40fe87..460c28e013 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -2,5 +2,14 @@
 /.git
 .gitattributes
 .gitignore
+.ipynb_checkpoints
+.pytest_cache
+.virtual_documents
+Untitled.*
 checkouts
 LICENSE
+node_modules
+Dockerfile
+coverage
+data_studio_jupyter_extensions.egg-info
+__pycache__
diff --git a/Dockerfile b/Dockerfile
index 38b6c53850..5e5a7cf05b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -49,11 +49,8 @@ RUN pip install -q -r requirements.txt
 # Install local extension package
 RUN mkdir extensions
 COPY . ./extensions
-RUN which pip
-RUN which python
 RUN pip install -U ./extensions
-RUN cat extensions/data_studio_jupyter_extensions/_version.py
-RUN python -c "import data_studio_jupyter_extensions; print(data_studio_jupyter_extensions.__version__)"
+RUN cd extensions && ./deployment/bin/version_check.sh
 
 #######################
 # Prepare runtime image
@@ -83,8 +80,7 @@ COPY --chown=${NB_USER}:${NB_USER} deployment/config/condarc ${APP_DIR}/.condarc
 COPY --chown=${NB_USER}:${NB_USER} deployment/config/npmrc ${APP_DIR}/.npmrc
 COPY --chown=${NB_USER}:${NB_USER} version ${APP_DIR}
 COPY --chown=${NB_USER}:${NB_USER} deployment ${APP_DIR}
-RUN python -c "import data_studio_jupyter_extensions; print(data_studio_jupyter_extensions.__version__)"
-RUN cat ${APP_DIR}/version
+RUN ${APP_DIR}/bin/version_check.sh
 
 USER ${NB_USER}
 WORKDIR ${APP_DIR}
diff --git a/deployment/bin/version_check.sh b/deployment/bin/version_check.sh
new file mode 100755
index 0000000000..590fbc1a93
--- /dev/null
+++ b/deployment/bin/version_check.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+version=$(cat version)
+echo "docker version: ${version}"
+cd ~
+py_version=$(python -c "import data_studio_jupyter_extensions; print(data_studio_jupyter_extensions.__version__)")
+echo "python version: ${py_version}"
+
+if [ $version != ${py_version} ]; then
+    echo "Version mismatch!"
+    exit 1
+fi
diff --git a/src/status.tsx b/src/status.tsx
index 6c0185a7a7..c5d4020430 100644
--- a/src/status.tsx
+++ b/src/status.tsx
@@ -29,14 +29,6 @@ import theme from '@tidbits/react-tidbits/theme';
 
 import { ReactWidget } from '@jupyterlab/apputils';
 
-// TODO: thread config through PageConfig instead of detecting env
-// once the functionality is available in JupyterServer
-declare var process: {
-  env: {
-    NODE_ENV: string;
-  };
-};
-
 // Define the error states
 // https://github.pie.apple.com/pie-data-studio/notebook-service/blob/761d63604966db5918d2e491c0f89cce454b7f67/app/com/apple/datastudio/model/ResourceState.scala#L20
 const ERROR_STATES = [
@@ -95,11 +87,10 @@ export const kernelStatusPlugin: JupyterFrontEndPlugin<void> = {
     console.log('JupyterLab extension "Kernel Status" is activated!');
 
     const pathCache = new LRUCache(50);
-    let url = URLExt.join(PageConfig.getWsUrl(), 'subscribe');
+    let url =
+      PageConfig.getOption('studioSubscribeURL') ||
+      URLExt.join(PageConfig.getWsUrl(), 'subscribe');
 
-    if (process.env.NODE_ENV == 'test') {
-      url = 'ws://localhost:5555';
-    }
     const ws = new WebSocket(url);
 
     const getPath = (kernel_id: string) => {
diff --git a/test/status.spec.ts b/test/status.spec.ts
index 09f8470f13..2bdcf7b171 100644
--- a/test/status.spec.ts
+++ b/test/status.spec.ts
@@ -24,6 +24,7 @@ import WS from 'jest-websocket-mock';
 import { kernelStatusPlugin } from '../src/status';
 
 import * as testUtils from './utils';
+import { PageConfig } from '@jupyterlab/coreutils';
 
 const server = new JupyterServer();
 const defaultRendermime = NBTestUtils.defaultRenderMime();
@@ -59,8 +60,9 @@ describe('status', () => {
       rendermime: defaultRendermime
     });
     app.docRegistry.addWidgetFactory(factory);
-
-    const server = new WS('ws://localhost:5555', { jsonProtocol: true });
+    const url = 'ws://localhost:5555';
+    PageConfig.setOption('studioSubscribeURL', url);
+    const server = new WS(url, { jsonProtocol: true });
     await kernelStatusPlugin.activate(app, loggerRegistry, notebookTracker);
 
     // Use the factory to create a panel and simulate the widget