From f9ee86a98a704f657ac25675411b39ebf00c839a Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Fri, 29 Sep 2023 12:24:22 +0000
Subject: [PATCH] test.py, exec_wrapper.py: improve error handling

---
 mmgen/exception.py      |  1 +
 scripts/exec_wrapper.py | 11 +++++++++--
 test/include/pexpect.py |  2 +-
 test/test.py            | 17 +++++++++--------
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/mmgen/exception.py b/mmgen/exception.py
index 96c47c74..51596b8b 100755
--- a/mmgen/exception.py
+++ b/mmgen/exception.py
@@ -51,6 +51,7 @@ class UserOptError(Exception):            mmcode = 1
 class NoLEDSupport(Exception):            mmcode = 1
 class MsgFileFailedSID(Exception):        mmcode = 1
 class TestSuiteException(Exception):      mmcode = 1
+class TestSuiteSpawnedScriptException(Exception): mmcode = 1
 
 # 2: yellow hl, message only
 class InvalidTokenAddress(Exception):     mmcode = 2
diff --git a/scripts/exec_wrapper.py b/scripts/exec_wrapper.py
index ef6e3654..c50a1748 100755
--- a/scripts/exec_wrapper.py
+++ b/scripts/exec_wrapper.py
@@ -40,11 +40,14 @@ def exec_wrapper_write_traceback(e,exit_val):
 
 	if os.getenv('EXEC_WRAPPER_TRACEBACK'):
 		import traceback
+		cwd = os.getcwd()
+		sys.path.insert(0,cwd)
+		from test.overlay import get_overlay_tree_dir
+		overlay_path_pfx = os.path.relpath(get_overlay_tree_dir(cwd)) + '/'
 
-		cwd = os.path.abspath('.')
 		def fixup_fn(fn_in):
 			from mmgen.util2 import removeprefix,removesuffix
-			fn = removeprefix(removeprefix(fn_in,cwd+'/'),'test/overlay/tree/')
+			fn = removeprefix(removeprefix(fn_in,cwd+'/'),overlay_path_pfx)
 			return removesuffix(fn,'_orig.py') + '.py' if fn.endswith('_orig.py') else fn
 			# Python 3.9:
 			# fn = fn_in.removeprefix(cwd+'/').removeprefix('test/overlay/tree/')
@@ -69,6 +72,10 @@ def exec_wrapper_write_traceback(e,exit_val):
 		from test.include.common import test_py_error_fn
 		with open(test_py_error_fn,'w') as fp:
 			fp.write('\n'.join(tb_lines + [exc_line]))
+
+		print(c.blue('{} script exited with error').format(
+			'Test' if os.path.dirname(sys.argv[0]) == 'test' else 'Spawned' ))
+
 	else:
 		sys.stdout.write( c.purple((f'NONZERO_EXIT[{exit_val}]: ' if exit_val else '') + exc_line) + '\n' )
 
diff --git a/test/include/pexpect.py b/test/include/pexpect.py
index bcafd313..14cba199 100755
--- a/test/include/pexpect.py
+++ b/test/include/pexpect.py
@@ -92,7 +92,7 @@ class MMGenPexpect:
 		self.p.read()
 		ret = self.p.wait()
 		if ret != self.req_exit_val and not cfg.coverage:
-			die(1,red(f'test.py: spawned program exited with value {ret}'))
+			die( 'TestSuiteSpawnedScriptException', f'Spawned script exited with value {ret}' )
 		if cfg.profile:
 			return
 		if not self.skip_ok:
diff --git a/test/test.py b/test/test.py
index 1f6ae28c..b4c6be5f 100755
--- a/test/test.py
+++ b/test/test.py
@@ -1001,7 +1001,7 @@ elif cmd_args and cmd_args[0] in utils:
 if cfg.pause:
 	set_restore_term_at_exit()
 
-from mmgen.exception import TestSuiteException,TestSuiteFatalException
+from mmgen.exception import TestSuiteException,TestSuiteFatalException,TestSuiteSpawnedScriptException
 
 try:
 	tr = TestSuiteRunner(data_dir,trash_dir)
@@ -1018,13 +1018,14 @@ except TestSuiteException as e:
 	die(2,e.args[0])
 except TestSuiteFatalException as e:
 	die(4,e.args[0])
+except TestSuiteSpawnedScriptException as e:
+	# if spawned script is not running under exec_wrapper, output brief error msg:
+	if os.getenv('MMGEN_EXEC_WRAPPER'):
+		Msg(red(str(e)))
+		Msg(blue('test.py: spawned script exited with error'))
 except Exception:
-	if os.getenv('MMGEN_EXEC_WRAPPER'): # test.py itself is running under exec_wrapper
-		import traceback
-		print(''.join(traceback.format_exception(*sys.exc_info())))
-		msg(blue('Test script exited with error'))
-	else:
-		msg(blue('Spawned script exited with error'))
-	raise
+	# if test.py itself is running under exec_wrapper, re-raise so exec_wrapper can handle exception:
+	if os.getenv('MMGEN_EXEC_WRAPPER'):
+		raise
 except:
 	raise