4 files changed, 7 insertions, 7 deletions
diff --git a/tests/micropython/heapalloc_exc_raise.py b/tests/micropython/heapalloc_exc_raise.py
index d60e14ce5..fb63a84bf 100644
--- a/tests/micropython/heapalloc_exc_raise.py
+++ b/tests/micropython/heapalloc_exc_raise.py
@@ -5,6 +5,7 @@ import micropython
 e = ValueError("error")
 
 def func():
+    micropython.heap_lock()
     try:
         # This works as is because traceback is not allocated
         # if not possible (heap is locked, no memory). If heap
@@ -16,8 +17,7 @@ def func():
         raise e
     except Exception as e2:
         print(e2)
+    micropython.heap_unlock()
 
-micropython.heap_lock()
 func()
 print("ok")
-micropython.heap_unlock()
diff --git a/tests/micropython/heapalloc_iter.py b/tests/micropython/heapalloc_iter.py
index 79461f999..30ac82e14 100644
--- a/tests/micropython/heapalloc_iter.py
+++ b/tests/micropython/heapalloc_iter.py
@@ -11,8 +11,10 @@ except (ImportError, AttributeError):
     heap_lock = heap_unlock = lambda:0
 
 def do_iter(l):
+    heap_lock()
     for i in l:
         print(i)
+    heap_unlock()
 
 def gen_func():
     yield 1
@@ -55,7 +57,6 @@ print(sum(t))
 heap_unlock()
 
 # test iterating over collections with the heap locked
-heap_lock()
 do_iter(b'123')
 do_iter(ba)
 do_iter(ar)
@@ -66,4 +67,3 @@ do_iter(s)
 do_iter(fs)
 do_iter(g1)
 do_iter(g2)
-heap_unlock()
diff --git a/tests/micropython/heapalloc_traceback.py b/tests/micropython/heapalloc_traceback.py
index f4212b6ce..813dea4b2 100644
--- a/tests/micropython/heapalloc_traceback.py
+++ b/tests/micropython/heapalloc_traceback.py
@@ -16,17 +16,17 @@ except:
     pass
 
 def test():
+    micropython.heap_lock()
     global global_exc
     global_exc.__traceback__ = None
     try:
         raise global_exc
     except StopIteration:
         print('StopIteration')
+    micropython.heap_unlock()
 
 # call test() with heap allocation disabled
-micropython.heap_lock()
 test()
-micropython.heap_unlock()
 
 # print the exception that was raised
 buf = uio.StringIO()
diff --git a/tests/micropython/heapalloc_traceback.py.exp b/tests/micropython/heapalloc_traceback.py.exp
index 291bbd697..facd0af13 100644
--- a/tests/micropython/heapalloc_traceback.py.exp
+++ b/tests/micropython/heapalloc_traceback.py.exp
@@ -1,5 +1,5 @@
 StopIteration
 Traceback (most recent call last):
-  File  , line 22, in test
+  File  , line 23, in test
 StopIteration: