diff --git a/opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_4_tests.cpp b/opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_4_tests.cpp
index fb75506565..b1cc0b8e23 100644
--- a/opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_4_tests.cpp
+++ b/opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_4_tests.cpp
@@ -647,6 +647,9 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenStaticPartitioningEnabledWhen
     mockCsr.storeMakeResidentAllocations = true;
 
     DispatchFlags dispatchFlags = DispatchFlagsHelper::createDefaultDispatchFlags();
+
+    cleanupHeaps();
+    initHeaps();
     mockCsr.flushTask(commandStream,
                       0,
                       &dsh,
@@ -759,4 +762,4 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenTagValueNotMeetingTaskCountTo
 
     const auto ret = mockCsr->waitForCompletionWithTimeout(WaitParams{true, true, 10}, taskCountToWait);
     EXPECT_EQ(NEO::WaitStatus::NotReady, ret);
-}
\ No newline at end of file
+}
diff --git a/opencl/test/unit_test/fixtures/ult_command_stream_receiver_fixture.h b/opencl/test/unit_test/fixtures/ult_command_stream_receiver_fixture.h
index 27a57128da..c71b41ec45 100644
--- a/opencl/test/unit_test/fixtures/ult_command_stream_receiver_fixture.h
+++ b/opencl/test/unit_test/fixtures/ult_command_stream_receiver_fixture.h
@@ -30,8 +30,6 @@ struct UltCommandStreamReceiverTest
         ClDeviceFixture::SetUp();
         ClHardwareParse::SetUp();
 
-        size_t sizeStream = 512;
-        size_t alignmentStream = 0x1000;
         cmdBuffer = alignedMalloc(sizeStream, alignmentStream);
         dshBuffer = alignedMalloc(sizeStream, alignmentStream);
         iohBuffer = alignedMalloc(sizeStream, alignmentStream);
@@ -42,6 +40,14 @@ struct UltCommandStreamReceiverTest
         ASSERT_NE(nullptr, iohBuffer);
         ASSERT_NE(nullptr, sshBuffer);
 
+        initHeaps();
+
+        flushTaskFlags.threadArbitrationPolicy = NEO::HwHelper::get(hardwareInfo.platform.eRenderCoreFamily).getDefaultThreadArbitrationPolicy();
+
+        pDevice->getGpgpuCommandStreamReceiver().setupContext(*pDevice->getDefaultEngine().osContext);
+    }
+
+    void initHeaps() {
         commandStream.replaceBuffer(cmdBuffer, sizeStream);
         auto graphicsAllocation = new MockGraphicsAllocation(cmdBuffer, sizeStream);
         commandStream.replaceGraphicsAllocation(graphicsAllocation);
@@ -58,18 +64,18 @@ struct UltCommandStreamReceiverTest
         ssh.replaceBuffer(sshBuffer, sizeStream);
         graphicsAllocation = new MockGraphicsAllocation(sshBuffer, sizeStream);
         ssh.replaceGraphicsAllocation(graphicsAllocation);
-
-        flushTaskFlags.threadArbitrationPolicy = NEO::HwHelper::get(hardwareInfo.platform.eRenderCoreFamily).getDefaultThreadArbitrationPolicy();
-
-        pDevice->getGpgpuCommandStreamReceiver().setupContext(*pDevice->getDefaultEngine().osContext);
     }
 
-    void TearDown() override {
-        pDevice->getGpgpuCommandStreamReceiver().flushBatchedSubmissions();
+    void cleanupHeaps() {
         delete dsh.getGraphicsAllocation();
         delete ioh.getGraphicsAllocation();
         delete ssh.getGraphicsAllocation();
         delete commandStream.getGraphicsAllocation();
+    }
+
+    void TearDown() override {
+        pDevice->getGpgpuCommandStreamReceiver().flushBatchedSubmissions();
+        cleanupHeaps();
 
         alignedFree(sshBuffer);
         alignedFree(iohBuffer);
@@ -167,5 +173,8 @@ struct UltCommandStreamReceiverTest
     uint32_t latestSentDcFlushTaskCount;
     uint32_t latestSentNonDcFlushTaskCount;
     uint32_t dcFlushRequiredTaskCount;
+
+    const size_t sizeStream = 512;
+    const size_t alignmentStream = 0x1000;
 };
 } // namespace NEO
diff --git a/opencl/test/unit_test/mocks/mock_buffer.h b/opencl/test/unit_test/mocks/mock_buffer.h
index da5556331a..aeaefdd8c4 100644
--- a/opencl/test/unit_test/mocks/mock_buffer.h
+++ b/opencl/test/unit_test/mocks/mock_buffer.h
@@ -20,9 +20,16 @@ class MockBufferStorage {
   public:
     MockBufferStorage() : mockGfxAllocation(data, sizeof(data) / 2),
                           multiGfxAllocation(GraphicsAllocationHelper::toMultiGraphicsAllocation(&mockGfxAllocation)) {
+        initDevice();
     }
+
     MockBufferStorage(bool unaligned) : mockGfxAllocation(unaligned ? alignUp(&data, 4) : alignUp(&data, 64), sizeof(data) / 2),
                                         multiGfxAllocation(GraphicsAllocationHelper::toMultiGraphicsAllocation(&mockGfxAllocation)) {
+        initDevice();
+    }
+    void initDevice() {
+        VariableBackup<uint32_t> maxOsContextCountBackup(&MemoryManager::maxOsContextCount);
+        device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
     }
     ~MockBufferStorage() {
         if (mockGfxAllocation.getDefaultGmm()) {
@@ -31,7 +38,7 @@ class MockBufferStorage {
     }
     char data[128];
     MockGraphicsAllocation mockGfxAllocation;
-    std::unique_ptr<MockDevice> device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
+    std::unique_ptr<MockDevice> device;
     MultiGraphicsAllocation multiGfxAllocation;
 };
 
diff --git a/shared/test/unit_test/command_container/command_container_tests.cpp b/shared/test/unit_test/command_container/command_container_tests.cpp
index 12081b6c6b..7706120422 100644
--- a/shared/test/unit_test/command_container/command_container_tests.cpp
+++ b/shared/test/unit_test/command_container/command_container_tests.cpp
@@ -616,10 +616,12 @@ TEST_P(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenHe
 
     auto executionEnvironment = new NEO::ExecutionEnvironment();
     const size_t numDevices = 2;
+
     executionEnvironment->prepareRootDeviceEnvironments(numDevices);
     for (auto i = 0u; i < numDevices; i++) {
         executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
     }
+    executionEnvironment->calculateMaxOsContextCount();
     auto device0 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 0u));
     auto device1 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 1u));
 
@@ -643,10 +645,12 @@ TEST_P(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenHe
 TEST_F(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenCmdBufferAllocationIsCreatedWithCorrectRootDeviceIndex) {
     auto executionEnvironment = new NEO::ExecutionEnvironment();
     const size_t numDevices = 2;
+
     executionEnvironment->prepareRootDeviceEnvironments(numDevices);
     for (auto i = 0u; i < numDevices; i++) {
         executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
     }
+    executionEnvironment->calculateMaxOsContextCount();
     auto device0 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 0u));
     auto device1 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 1u));
 
@@ -670,6 +674,7 @@ TEST_F(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenIn
     for (auto i = 0u; i < numDevices; i++) {
         executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
     }
+    executionEnvironment->calculateMaxOsContextCount();
     auto device0 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 0u));
     auto device1 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 1u));
 
@@ -783,4 +788,4 @@ TEST_F(CommandContainerTest, givenCmdContainerWhenCloseAndAllocateNextCommandBuf
     EXPECT_EQ(cmdContainer.getCmdBufferAllocations().size(), 1u);
     cmdContainer.closeAndAllocateNextCommandBuffer();
     EXPECT_EQ(cmdContainer.getCmdBufferAllocations().size(), 2u);
-}
\ No newline at end of file
+}