diff --git a/vm/src/vm.rs b/vm/src/vm.rs
index a1c7e5868..c8a131552 100644
--- a/vm/src/vm.rs
+++ b/vm/src/vm.rs
@@ -303,33 +303,28 @@ impl VirtualMachine {
     }
 
     pub fn import(&self, module: &str, from_list: &PyObjectRef, level: usize) -> PyResult {
-        let sys_modules = self
-            .get_attribute(self.sys_module.clone(), "modules")
-            .unwrap();
-        if let Ok(module) = sys_modules.get_item(module.to_string(), self) {
-            Ok(module)
-        } else {
-            match self.get_attribute(self.builtins.clone(), "__import__") {
-                Ok(func) => self.invoke(
-                    func,
-                    vec![
-                        self.ctx.new_str(module.to_string()),
-                        if self.current_frame().is_some() {
-                            self.get_locals().into_object()
-                        } else {
-                            self.get_none()
-                        },
-                        self.get_none(),
-                        from_list.clone(),
-                        self.ctx.new_int(level),
-                    ],
-                ),
-                Err(_) => Err(self.new_exception(
-                    self.ctx.exceptions.import_error.clone(),
-                    "__import__ not found".to_string(),
-                )),
-            }
-        }
+        let sys_modules = self.get_attribute(self.sys_module.clone(), "modules")?;
+        sys_modules.get_item(module.to_string(), self).or_else(|_| {
+            let import_func = self
+                .get_attribute(self.builtins.clone(), "__import__")
+                .map_err(|_| self.new_import_error("__import__ not found".to_string()))?;
+
+            let locals = if let Some(frame) = self.current_frame() {
+                frame.scope.get_locals().into_object()
+            } else {
+                self.get_none()
+            };
+            self.invoke(
+                import_func,
+                vec![
+                    self.ctx.new_str(module.to_string()),
+                    locals,
+                    self.get_none(),
+                    from_list.clone(),
+                    self.ctx.new_int(level),
+                ],
+            )
+        })
     }
 
     /// Determines if `obj` is an instance of `cls`, either directly, indirectly or virtually via
diff --git a/wasm/demo/package.json b/wasm/demo/package.json
index a03c80ace..6f66b120d 100644
--- a/wasm/demo/package.json
+++ b/wasm/demo/package.json
@@ -5,6 +5,7 @@
     "main": "index.js",
     "dependencies": {
         "codemirror": "^5.42.0",
+        "serve": "^11.0.2",
         "xterm": "^3.8.0"
     },
     "devDependencies": {
@@ -24,7 +25,7 @@
         "build": "webpack",
         "dist": "webpack --mode production",
         "test": "cd ../tests; pipenv run pytest",
-        "ci": "start-server-and-test dev http-get://localhost:8080 test"
+        "ci": "start-server-and-test 'npm run build && serve dist/ -n -p 8080' :8080 test"
     },
     "repository": {
         "type": "git",
diff --git a/wasm/tests/.gitignore b/wasm/tests/.gitignore
new file mode 100644
index 000000000..283f43cb0
--- /dev/null
+++ b/wasm/tests/.gitignore
@@ -0,0 +1 @@
+geckodriver.log
diff --git a/wasm/tests/test_demo.py b/wasm/tests/test_demo.py
index 9dce40681..a9694f071 100644
--- a/wasm/tests/test_demo.py
+++ b/wasm/tests/test_demo.py
@@ -1,4 +1,5 @@
 import time
+import sys
 
 from selenium import webdriver
 from selenium.webdriver.firefox.options import Options
@@ -16,12 +17,23 @@ vm.destroy();
 return output;
 """
 
+def print_stack(driver):
+	stack = driver.execute_script(
+		"return window.__RUSTPYTHON_ERROR_MSG + '\\n' + window.__RUSTPYTHON_ERROR_STACK"
+	)
+	print(f"RustPython error stack:\n{stack}", file=sys.stderr)
+
+
 @pytest.fixture(scope="module")
 def driver(request):
 	options = Options()
 	options.add_argument('-headless')
 	driver = webdriver.Firefox(options=options)
-	driver.get("http://localhost:8080")
+	try:
+		driver.get("http://localhost:8080")
+	except Exception as e:
+		print_stack(driver)
+		raise
 	time.sleep(5)
 	yield driver
 	driver.close()
@@ -35,4 +47,9 @@ def driver(request):
 )
 def test_demo(driver, script, output):
 	script = RUN_CODE_TEMPLATE.format(script)
-	assert driver.execute_script(script).strip() == output
+	try:
+		script_output = driver.execute_script(script)
+	except Exception as e:
+		print_stack(driver)
+		raise
+	assert script_output.strip() == output