diff --git a/gdb-port/tests/integration/test_malloc_attribution.py b/gdb-port/tests/integration/test_malloc_attribution.py
index 3c3210edbdfc550209e206b587674e19e1fd6e35..4460f41fedb4676044a55813db5262a5cec2692c 100644
--- a/gdb-port/tests/integration/test_malloc_attribution.py
+++ b/gdb-port/tests/integration/test_malloc_attribution.py
@@ -10,7 +10,7 @@ class TestMallocParserAttribution(unittest.TestCase):
 		# TODO: the above two lines simulate a side effect of enter_h_do_parse. could be avoided if self.a and self.b was assigned after the calls below
 		self.top_level_parse.enter_h_packrat_parse(self.a.address)
 
-#TODO: are we initializing from scratch for every test?
+		# TODO: are we initializing from scratch for every test?
 		self.top_level_parse.enter_h_do_parse(400, 256, self.a.address)
 		self.top_level_parse.parse_virtual(self.a.address)
 		self.top_level_parse.enter_perform_lowlevel_parse(self.a.address)
@@ -19,17 +19,47 @@ class TestMallocParserAttribution(unittest.TestCase):
 	def test_alloc_after_h_do_parse(self):
 		self.top_level_parse.enter_h_do_parse(400, 256, self.b.address)
 		self.top_level_parse.enter_h_arena_malloc_raw(150)
-		self.assertEqual(b.bytes_used[256], 150)
+		self.assertEqual(self.b.bytes_used[256], 150)
 
 	def test_alloc_after_parse_virtual(self):
 		self.top_level_parse.enter_h_do_parse(400, 256, self.b.address)
 		self.top_level_parse.parse_virtual(self.b.address)
 		self.top_level_parse.enter_h_arena_malloc_raw(150)
-		self.assertEqual(b.bytes_used[256], 150)
+		self.assertEqual(self.b.bytes_used[256], 150)
 
 	def test_alloc_after_enter_perform_lowlevel_parse(self):
 		self.top_level_parse.enter_h_do_parse(400, 256, self.b.address)
 		self.top_level_parse.parse_virtual(self.b.address)
 		self.top_level_parse.enter_perform_lowlevel_parse(self.b.address)
 		self.top_level_parse.enter_h_arena_malloc_raw(150)
-		self.assertEqual(b.bytes_used[256], 150)
+		self.assertEqual(self.b.bytes_used[256], 150)
+
+class TestMallocParserAttributionOnReturns(unittest.TestCase):
+	def setUp(self):
+		self.a = Parser("a", 128)
+		self.b = Parser("b", 32)
+		self.top_level_parse = TopLevelParse()
+		self.top_level_parse.parser_objs[32] = self.b
+		self.top_level_parse.parser_objs[128] = self.a
+		# TODO: the above two lines simulate a side effect of enter_h_do_parse. could be avoided if self.a and self.b was assigned after the calls below
+		self.top_level_parse.enter_h_packrat_parse(self.a.address)
+
+		# TODO: are we initializing from scratch for every test?
+		self.top_level_parse.enter_h_do_parse(400, 256, self.a.address)
+		self.top_level_parse.parse_virtual(self.a.address)
+		self.top_level_parse.enter_perform_lowlevel_parse(self.a.address)
+
+	def test_alloc_after_return_from_perform_lowlevel_parse(self):
+		self.top_level_parse.return_from_perform_lowlevel_parse()
+		self.top_level_parse.enter_h_arena_malloc_raw(150)
+		self.assertEqual(self.b.bytes_used[256], 150)
+
+	def test_alloc_after_return_from_h_do_parse(self):
+		hps_t = gdb.lookup_type("HParseState")
+		test_parse_state = gdb.Value(b'\x00'*80, hps_t)
+
+		self.top_level_parse.return_from_perform_lowlevel_parse()
+		self.top_level_parse.return_from_h_do_parse(test_parse_state, self.b.address, 0)
+		self.top_level_parse.enter_h_arena_malloc_raw(150)
+		self.assertEqual(self.a.bytes_used[256], 150)
+