Fixed exception handling for vizviewer (#354)

gaogaotiantian · Jul 17, 2023 · 7cc2e5c · 7cc2e5c
1 parent 67219a9
commit 7cc2e5c
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/src/viztracer/viewer.py b/src/viztracer/viewer.py
@@ -14,6 +14,7 @@
 import socketserver
 import subprocess
 import sys
+import traceback
 import threading
 import time
 import urllib.parse
@@ -298,9 +299,14 @@ def __init__(
         super().__init__(daemon=True)
 
     def run(self) -> None:
-        self.retcode = self.view()
-        # If it returns from view(), also set ready
-        self.ready.set()
+        try:
+            self.retcode = self.view()
+        except Exception:
+            self.retcode = 1
+            traceback.print_exc()
+        finally:
+            # If it returns from view(), also set ready
+            self.ready.set()
 
     def view(self) -> int:
         # Get file data

diff --git a/tests/cmdline_tmpl.py b/tests/cmdline_tmpl.py
@@ -118,20 +118,21 @@ def template(self,
                 logging.error(f"stdout: {e.stdout}")
                 logging.error(f"stderr: {e.stderr}")
                 raise e
-        if not (success ^ (result.returncode != 0)):
+        expected = (success ^ (result.returncode != 0))
+        if not expected:
             logging.error(f"return code: {result.returncode}")
             logging.error(f"stdout:\n{result.stdout.decode('utf-8')}")
             logging.error(f"stderr:\n{result.stderr.decode('utf-8')}")
-        self.assertTrue(success ^ (result.returncode != 0))
-        if success:
-            if expected_output_file:
+        self.assertTrue(expected)
+        if expected:
+            if success and expected_output_file:
                 if type(expected_output_file) is list:
                     for f in expected_output_file:
                         self.assertFileExists(f)
                 elif type(expected_output_file) is str:
                     self.assertFileExists(expected_output_file)
 
-            if expected_entries:
+            if success and expected_entries:
                 assert (type(expected_output_file) is str and expected_output_file.split(".")[-1] == "json")
                 with open(expected_output_file) as f:
                     data = json.load(f)

diff --git a/tests/test_viewer.py b/tests/test_viewer.py
@@ -457,6 +457,11 @@ def test_directory_max_port(self):
         finally:
             shutil.rmtree(tmp_dir)
 
+    def test_exception(self):
+        test_data_dir = os.path.join(os.path.dirname(__file__), "data")
+        self.template(["vizviewer", "--port", "-3", os.path.join(test_data_dir, "fib.json")],
+                      success=False, expected_output_file=None, expected_stderr=".*Traceback.*")
+
     def test_invalid(self):
         self.template(["vizviewer", "do_not_exist.json"], success=False, expected_output_file=None)
         self.template(["vizviewer", "README.md"], success=False, expected_output_file=None)