diff --git a/shared/test/common/os_interface/linux/sys_calls_linux_ult.cpp b/shared/test/common/os_interface/linux/sys_calls_linux_ult.cpp
index 6262cd3a4a..334edbd83d 100644
--- a/shared/test/common/os_interface/linux/sys_calls_linux_ult.cpp
+++ b/shared/test/common/os_interface/linux/sys_calls_linux_ult.cpp
@@ -74,6 +74,7 @@ bool mmapAllowExtendedPointers = false;
 bool failMmap = false;
 uint32_t mmapFuncCalled = 0u;
 uint32_t munmapFuncCalled = 0u;
+bool failMunmap = false;
 
 int (*sysCallsOpen)(const char *pathname, int flags) = nullptr;
 int (*sysCallsClose)(int fileDescriptor) = nullptr;
@@ -324,6 +325,10 @@ void *mmap(void *addr, size_t size, int prot, int flags, int fd, off_t off) noex
 
 int munmap(void *addr, size_t size) noexcept {
     munmapFuncCalled++;
+    if (failMunmap) {
+        return -1;
+    }
+
     auto ptrIt = std::find(mmapVector.begin(), mmapVector.end(), addr);
     if (ptrIt != mmapVector.end()) {
         mmapVector.erase(ptrIt);
diff --git a/shared/test/common/os_interface/linux/xe/mock_drm_xe.inl b/shared/test/common/os_interface/linux/xe/mock_drm_xe.inl
index 237af07ab5..db330980c3 100644
--- a/shared/test/common/os_interface/linux/xe/mock_drm_xe.inl
+++ b/shared/test/common/os_interface/linux/xe/mock_drm_xe.inl
@@ -27,6 +27,8 @@ struct DrmMockXe : public DrmMockCustom {
     int gemVmBindReturn = 0;
     GemClose passedGemClose{};
     int gemCloseCalled = 0;
+    int gemMmapOffsetCalled = 0;
+    int gemDestroyContextCalled = 0;
 
     const uint16_t revId = 0x12;
     const uint16_t devId = 0xabc;
@@ -58,6 +60,7 @@ struct DrmMockXe : public DrmMockCustom {
     uint16_t createParamsCpuCaching = 0u;
     uint32_t createParamsPlacement = 0u;
     bool ioctlCalled = false;
+    bool forceMmapOffsetFail = false;
 
   protected:
     // Don't call directly, use the create() function
diff --git a/shared/test/common/os_interface/linux/xe/mock_drm_xe_definitions.inl b/shared/test/common/os_interface/linux/xe/mock_drm_xe_definitions.inl
index d771d7949c..0f0f419d23 100644
--- a/shared/test/common/os_interface/linux/xe/mock_drm_xe_definitions.inl
+++ b/shared/test/common/os_interface/linux/xe/mock_drm_xe_definitions.inl
@@ -47,11 +47,16 @@ int DrmMockXe::ioctl(DrmIoctl request, void *arg) {
             ret = 0;
     } break;
     case DrmIoctl::gemMmapOffset: {
-        struct drm_xe_gem_mmap_offset *v = static_cast<struct drm_xe_gem_mmap_offset *>(arg);
-        if (v->handle == testValueMapOff) {
-            v->offset = v->handle;
-            ret = 0;
+        gemMmapOffsetCalled++;
+
+        if (forceMmapOffsetFail) {
+            ret = -1;
+            break;
         }
+        struct drm_xe_gem_mmap_offset *v = static_cast<struct drm_xe_gem_mmap_offset *>(arg);
+        v->offset = v->handle;
+        ret = 0;
+
     } break;
     case DrmIoctl::primeFdToHandle: {
         PrimeHandle *v = static_cast<PrimeHandle *>(arg);
@@ -174,6 +179,7 @@ int DrmMockXe::ioctl(DrmIoctl request, void *arg) {
         ret = 0;
     } break;
     case DrmIoctl::gemContextDestroy: {
+        gemDestroyContextCalled++;
         ret = 0;
     } break;
     case DrmIoctl::perfOpen: {