diff --git a/src/backend/storage/file/test/compress_zlib_test.c b/src/backend/storage/file/test/compress_zlib_test.c
index f94275c474e421ca7668b08a156404c9bdb2c651..68d80b28514b9400a41d7f83a2ea791e05880c51 100644
--- a/src/backend/storage/file/test/compress_zlib_test.c
+++ b/src/backend/storage/file/test/compress_zlib_test.c
@@ -80,11 +80,13 @@ test__bfz_zlib_init__palloc_write(void **state)
 	bfz.mode = BFZ_MODE_APPEND;
 	bfz.fd = -1;
 
-	Size beforeAlloc = MemoryContextGetPeakSpace(TopMemoryContext);
+	Size beforeAlloc = MemoryContextGetPeakSpace(TopTransactionContext);
+	assert_true(TopTransactionContext != CurrentMemoryContext);
 
 	bfz_zlib_init(&bfz);
 
-	Size afterAlloc = MemoryContextGetPeakSpace(TopMemoryContext);
+	Size afterAlloc = MemoryContextGetPeakSpace(TopTransactionContext);
+	assert_true(TopTransactionContext != CurrentMemoryContext);
 
 	int memZlib = zlib_memory_needed(true /* isWrite */);
 
@@ -100,11 +102,13 @@ test__bfz_zlib_init__palloc_read(void **state)
 	bfz.mode = BFZ_MODE_SCAN;
 	bfz.fd = -1;
 
-	Size beforeAlloc = MemoryContextGetPeakSpace(TopMemoryContext);
+	Size beforeAlloc = MemoryContextGetPeakSpace(TopTransactionContext);
+	assert_true(TopTransactionContext != CurrentMemoryContext);
 
 	bfz_zlib_init(&bfz);
 
-	Size afterAlloc = MemoryContextGetPeakSpace(TopMemoryContext);
+	Size afterAlloc = MemoryContextGetPeakSpace(TopTransactionContext);
+	assert_true(TopTransactionContext != CurrentMemoryContext);
 
 	int memZlib = zlib_memory_needed(false /* isWrite */);
 
@@ -116,6 +120,13 @@ main(int argc, char* argv[])
 {
 	cmockery_parse_arguments(argc, argv);
 
+	TopTransactionContext =
+		AllocSetContextCreate(TopMemoryContext,
+							  "TopTransactionContext",
+							  ALLOCSET_DEFAULT_MINSIZE,
+							  ALLOCSET_DEFAULT_INITSIZE,
+							  ALLOCSET_DEFAULT_MAXSIZE);
+
 	const UnitTest tests[] = {
 		unit_test(test__bfz_zlib_init__palloc_write),
 		unit_test(test__bfz_zlib_init__palloc_read)