2 files changed, 17 insertions, 3 deletions
diff --git a/tests/micropython/heapalloc_traceback.py b/tests/micropython/heapalloc_traceback.py
index 9f43da2592..808df0225c 100644
--- a/tests/micropython/heapalloc_traceback.py
+++ b/tests/micropython/heapalloc_traceback.py
@@ -2,6 +2,7 @@
 
 import micropython
 import sys
+import uio
 
 # preallocate exception instance with some room for a traceback
 global_exc = StopIteration()
@@ -15,10 +16,21 @@ def test():
     global_exc.__traceback__ = None
     try:
         raise global_exc
-    except StopIteration as e:
-        sys.print_exception(e)
+    except StopIteration:
+        print('StopIteration')
 
 # call test() with heap allocation disabled
 micropython.heap_lock()
 test()
 micropython.heap_unlock()
+
+# print the exception that was raised
+buf = uio.StringIO()
+sys.print_exception(global_exc, buf)
+for l in buf.getvalue().split("\n"):
+    # uPy on pyboard prints <stdin> as file, so remove filename.
+    if l.startswith("  File "):
+        l = l.split('"')
+        print(l[0], l[2])
+    else:
+        print(l)
diff --git a/tests/micropython/heapalloc_traceback.py.exp b/tests/micropython/heapalloc_traceback.py.exp
index 16284f0da3..9ba10948d5 100644
--- a/tests/micropython/heapalloc_traceback.py.exp
+++ b/tests/micropython/heapalloc_traceback.py.exp
@@ -1,3 +1,5 @@
+StopIteration
 Traceback (most recent call last):
-  File "micropython/heapalloc_traceback.py", line 17, in test
+  File  , line 18, in test
 StopIteration: 
+