/*
 * Copyright (C) 2020-2026 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#include "shared/source/gmm_helper/gmm.h"
#include "shared/source/memory_manager/gfx_partition.h"
#include "shared/source/os_interface/device_factory.h"
#include "shared/source/unified_memory/unified_memory.h"
#include "shared/source/utilities/cpu_info.h"
#include "shared/test/common/mocks/mock_bindless_heaps_helper.h"
#include "shared/test/common/mocks/mock_command_stream_receiver.h"
#include "shared/test/common/mocks/mock_cpu_page_fault_manager.h"
#include "shared/test/common/mocks/mock_device.h"
#include "shared/test/common/mocks/mock_graphics_allocation.h"
#include "shared/test/common/mocks/mock_memory_manager.h"
#include "shared/test/common/mocks/mock_memory_operations_handler.h"
#include "shared/test/common/mocks/mock_svm_manager.h"
#include "shared/test/common/test_macros/hw_test.h"

#include "level_zero/core/source/cmdqueue/cmdqueue.h"
#include "level_zero/core/source/context/context_imp.h"
#include "level_zero/core/source/driver/driver_handle.h"
#include "level_zero/core/source/image/image.h"
#include "level_zero/core/test/common/ult_helpers_l0.h"
#include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
#include "level_zero/core/test/unit_tests/fixtures/host_pointer_manager_fixture.h"
#include "level_zero/core/test/unit_tests/mocks/mock_cmdlist.h"
#include "level_zero/core/test/unit_tests/mocks/mock_context.h"
#include "level_zero/include/level_zero/driver_experimental/zex_context.h"

#include "gtest/gtest.h"

namespace L0 {
template <GFXCORE_FAMILY gfxCoreFamily>
struct CommandListCoreFamily;

namespace ult {

using MultiDeviceContextTests = Test<MultiDeviceFixture>;

TEST_F(MultiDeviceContextTests,
       whenCreatingContextWithZeroNumDevicesThenAllDevicesAreAssociatedWithTheContext) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(Context::fromHandle(hContext));

    for (size_t i = 0; i < driverHandle->devices.size(); i++) {
        EXPECT_NE(contextImp->getDevices().find(driverHandle->devices[i]->getRootDeviceIndex()), contextImp->getDevices().end());
    }

    res = L0::Context::fromHandle(hContext)->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(MultiDeviceContextTests,
       whenCreatingContextWithNonZeroNumDevicesThenOnlySpecifiedDeviceAndItsSubDevicesAreAssociatedWithTheContext) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_device_handle_t device0 = driverHandle->devices[0]->toHandle();
    Device *l0Device0 = static_cast<Device *>(device0);
    uint32_t subDeviceCount0 = 0;
    ze_result_t res = l0Device0->getSubDevices(&subDeviceCount0, nullptr);
    EXPECT_EQ(res, ZE_RESULT_SUCCESS);
    EXPECT_EQ(subDeviceCount0, numSubDevices);
    std::vector<ze_device_handle_t> subDevices0(subDeviceCount0);
    res = l0Device0->getSubDevices(&subDeviceCount0, subDevices0.data());
    EXPECT_EQ(res, ZE_RESULT_SUCCESS);

    ze_device_handle_t device1 = driverHandle->devices[1]->toHandle();
    Device *l0Device1 = static_cast<Device *>(device1);
    uint32_t subDeviceCount1 = 0;
    res = l0Device1->getSubDevices(&subDeviceCount1, nullptr);
    EXPECT_EQ(res, ZE_RESULT_SUCCESS);
    EXPECT_EQ(subDeviceCount1, numSubDevices);
    std::vector<ze_device_handle_t> subDevices1(subDeviceCount1);
    res = l0Device1->getSubDevices(&subDeviceCount1, subDevices1.data());
    EXPECT_EQ(res, ZE_RESULT_SUCCESS);

    uint32_t subSubDeviceCount1 = 0;
    res = static_cast<Device *>(subDevices1[0])->getSubDevices(&subSubDeviceCount1, nullptr);
    EXPECT_EQ(res, ZE_RESULT_SUCCESS);

    res = driverHandle->createContext(&desc, 1u, &device1, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(Context::fromHandle(hContext));

    uint32_t expectedDeviceCountInContext = 1;
    EXPECT_EQ(contextImp->getDevices().size(), expectedDeviceCountInContext);

    EXPECT_FALSE(contextImp->isDeviceDefinedForThisContext(L0::Device::fromHandle(device0)));
    for (auto subDevice : subDevices0) {
        EXPECT_FALSE(contextImp->isDeviceDefinedForThisContext(L0::Device::fromHandle(subDevice)));
    }

    EXPECT_TRUE(contextImp->isDeviceDefinedForThisContext(L0::Device::fromHandle(device1)));
    for (auto subDevice : subDevices1) {
        auto l0SubDevice = static_cast<Device *>(subDevice);
        EXPECT_TRUE(contextImp->isDeviceDefinedForThisContext(l0SubDevice));

        for (auto &subSubDevice : l0SubDevice->subDevices) {
            EXPECT_TRUE(contextImp->isDeviceDefinedForThisContext(L0::Device::fromHandle(subSubDevice)));
        }
    }

    res = L0::Context::fromHandle(hContext)->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(MultiDeviceContextTests,
       whenAllocatingDeviceMemoryWithDeviceNotDefinedForContextThenDeviceLostIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_device_handle_t device = driverHandle->devices[1]->toHandle();

    ze_result_t res = driverHandle->createContext(&desc, 1u, &device, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(Context::fromHandle(hContext));

    ze_device_mem_alloc_desc_t deviceDesc = {};
    size_t size = 4096;
    void *ptr = nullptr;
    res = contextImp->allocDeviceMem(driverHandle->devices[0]->toHandle(), &deviceDesc, size, 0u, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_DEVICE_LOST, res);

    res = L0::Context::fromHandle(hContext)->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(MultiDeviceContextTests,
       whenPeerAllocationForReservedAllocationIsRequestedThenPeerAllocationIsAddedToDeviceMapAndRemovedWhenAllocationIsFreed) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->devices[0]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();
    driverHandle->devices[1]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[1]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(driverHandle->devices[0], size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(driverHandle->devices[0], &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    uintptr_t peerGpuAddress = 0u;
    auto allocData = driverHandle->getSvmAllocsManager()->getSVMAlloc(ptr);
    EXPECT_NE(allocData, nullptr);
    EXPECT_EQ(driverHandle->getSvmAllocsManager()->allocationsCounter.load(), allocData->getAllocId());
    auto peerAlloc = driverHandle->getPeerAllocation(driverHandle->devices[1], allocData, ptr, &peerGpuAddress, nullptr);
    EXPECT_NE(peerAlloc, nullptr);

    Device *l0Device1 = static_cast<Device *>(driverHandle->devices[1]);

    {
        auto iter = l0Device1->peerAllocations.allocations.find(ptr);
        EXPECT_NE(iter, l0Device1->peerAllocations.allocations.end());
    }

    res = contextImp->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto iter = l0Device1->peerAllocations.allocations.find(ptr);
        EXPECT_EQ(iter, l0Device1->peerAllocations.allocations.end());
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(MultiDeviceContextTests,
       GivenDeviceMemoryWhenMakeResidentCalledOnPeerDeviceThenSuccessReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->devices[0]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();
    driverHandle->devices[1]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[1]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    const size_t size = 4096;
    void *ptr = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {};
    res = context->allocDeviceMem(driverHandle->devices[0],
                                  &deviceDesc,
                                  size,
                                  0,
                                  &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->makeMemoryResident(driverHandle->devices[1], ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeMem(ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(MultiDeviceContextTests,
       GivenInvalidDeviceMemoryWhenMakeResidentCalledOnPeerDeviceThenSuccessReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->devices[0]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();
    driverHandle->devices[1]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[1]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    const size_t size = 4096;
    void *ptr = reinterpret_cast<void *>(0x1234);

    res = contextImp->makeMemoryResident(driverHandle->devices[1], ptr, size);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(MultiDeviceContextTests,
       whenMappingReservedMemoryOnPhysicalMemoryOnMultiDeviceThenSuccessReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->devices[0]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();
    driverHandle->devices[1]->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[1]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(driverHandle->devices[0], size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize * 2, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(driverHandle->devices[0], &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    ze_physical_mem_handle_t secondHalfMem = {};
    res = contextImp->createPhysicalMem(driverHandle->devices[1], &descMem, &secondHalfMem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    std::vector<ze_memory_access_attribute_t> memoryAccessFlags = {
        ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE, ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY,
        ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};

    void *offsetAddr =
        reinterpret_cast<void *>(reinterpret_cast<uint64_t>(ptr) + pagesize);

    for (auto accessFlags : memoryAccessFlags) {
        res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->mapVirtualMem(offsetAddr, pagesize, secondHalfMem, offset, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->setVirtualMemAccessAttribute(ptr, pagesize, access);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->setVirtualMemAccessAttribute(offsetAddr, pagesize, access);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);

        ze_memory_access_attribute_t outAccess = {};
        size_t outSize = 0;
        res = contextImp->getVirtualMemAccessAttribute(ptr, pagesize, &outAccess, &outSize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        EXPECT_EQ(pagesize * 2, outSize);
        res = contextImp->getVirtualMemAccessAttribute(offsetAddr, pagesize, &outAccess, &outSize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        EXPECT_EQ(pagesize * 2, outSize);

        res = contextImp->unMapVirtualMem(ptr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->unMapVirtualMem(offsetAddr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    }

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroyPhysicalMem(secondHalfMem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize * 2);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(MultiDeviceContextTests,
       whenAllocatingSharedMemoryWithDeviceNotDefinedForContextThenDeviceLostIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_device_handle_t device = driverHandle->devices[1]->toHandle();

    ze_result_t res = driverHandle->createContext(&desc, 1u, &device, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(Context::fromHandle(hContext));

    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_host_mem_alloc_desc_t hostDesc = {};
    size_t size = 4096;
    void *ptr = nullptr;
    res = contextImp->allocSharedMem(driverHandle->devices[0]->toHandle(), &deviceDesc, &hostDesc, size, 0u, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_DEVICE_LOST, res);

    res = L0::Context::fromHandle(hContext)->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

struct SVMAllocsManagerContextMock : public NEO::SVMAllocsManager {
    SVMAllocsManagerContextMock(MemoryManager *memoryManager) : NEO::SVMAllocsManager(memoryManager) {}
    void *createHostUnifiedMemoryAllocation(size_t size, const UnifiedMemoryProperties &memoryProperties) override {
        EXPECT_EQ(expectedRootDeviceIndexes.size(), memoryProperties.rootDeviceIndices.size());
        EXPECT_NE(std::find(memoryProperties.rootDeviceIndices.begin(), memoryProperties.rootDeviceIndices.end(), expectedRootDeviceIndexes[0]),
                  memoryProperties.rootDeviceIndices.end());
        EXPECT_NE(std::find(memoryProperties.rootDeviceIndices.begin(), memoryProperties.rootDeviceIndices.end(), expectedRootDeviceIndexes[1]),
                  memoryProperties.rootDeviceIndices.end());
        return NEO::SVMAllocsManager::createHostUnifiedMemoryAllocation(size, memoryProperties);
    }

    std::vector<uint32_t> expectedRootDeviceIndexes;
};

using ContextGetStatusTest = Test<DeviceFixture>;
TEST_F(ContextGetStatusTest, givenCallToContextGetStatusThenCorrectErrorCodeIsReturnedWhenResourcesHaveBeenReleased) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    L0::Context *context = L0::Context::fromHandle(hContext);

    res = context->getStatus();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    L0UltHelper::cleanupUsmAllocPoolsAndReuse(driverHandle.get());

    for (auto device : driverHandle->devices) {
        L0::Device *l0Device = static_cast<Device *>(device);
        l0Device->releaseResources();
    }

    res = context->getStatus();
    EXPECT_EQ(ZE_RESULT_ERROR_DEVICE_LOST, res);
    context->destroy();
}

using ContextPowerSavingHintTest = Test<DeviceFixture>;
TEST_F(ContextPowerSavingHintTest, givenCallToContextCreateWithPowerHintDescThenPowerHintSetInDriverHandle) {
    ze_context_handle_t hContext;
    ze_context_desc_t ctxtDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC};
    ze_context_power_saving_hint_exp_desc_t powerHintContext = {};
    powerHintContext.stype = ZE_STRUCTURE_TYPE_POWER_SAVING_HINT_EXP_DESC;
    powerHintContext.hint = 1;
    ctxtDesc.pNext = &powerHintContext;
    ze_result_t res = driverHandle->createContext(&ctxtDesc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(powerHintContext.hint, driverHandle->powerHint);
    L0::Context *context = L0::Context::fromHandle(hContext);
    context->destroy();
}

TEST_F(ContextPowerSavingHintTest, givenCallToContextCreateWithPowerHintMinimumThenPowerHintSetInDriverHandle) {
    ze_context_handle_t hContext;
    ze_context_desc_t ctxtDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC};
    ze_context_power_saving_hint_exp_desc_t powerHintContext = {};
    powerHintContext.stype = ZE_STRUCTURE_TYPE_POWER_SAVING_HINT_EXP_DESC;
    powerHintContext.hint = ZE_POWER_SAVING_HINT_TYPE_MIN;
    ctxtDesc.pNext = &powerHintContext;
    ze_result_t res = driverHandle->createContext(&ctxtDesc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(powerHintContext.hint, driverHandle->powerHint);
    L0::Context *context = L0::Context::fromHandle(hContext);
    context->destroy();
}

TEST_F(ContextPowerSavingHintTest, givenCallToContextCreateWithPowerHintMaximumThenPowerHintSetInDriverHandle) {
    ze_context_handle_t hContext;
    ze_context_desc_t ctxtDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC};
    ze_context_power_saving_hint_exp_desc_t powerHintContext = {};
    powerHintContext.stype = ZE_STRUCTURE_TYPE_POWER_SAVING_HINT_EXP_DESC;
    powerHintContext.hint = ZE_POWER_SAVING_HINT_TYPE_MAX;
    ctxtDesc.pNext = &powerHintContext;
    ze_result_t res = driverHandle->createContext(&ctxtDesc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(powerHintContext.hint, driverHandle->powerHint);
    L0::Context *context = L0::Context::fromHandle(hContext);
    context->destroy();
}

TEST_F(ContextPowerSavingHintTest, givenCallToContextCreateWithPowerHintGreaterThanMaxHintThenErrorIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t ctxtDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_context_power_saving_hint_exp_desc_t powerHintContext = {};
    powerHintContext.stype = ZE_STRUCTURE_TYPE_POWER_SAVING_HINT_EXP_DESC;
    powerHintContext.hint = ZE_POWER_SAVING_HINT_TYPE_MAX + 1;
    ctxtDesc.pNext = &powerHintContext;
    ze_result_t res = driverHandle->createContext(&ctxtDesc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ENUMERATION, res);
}

TEST_F(ContextPowerSavingHintTest, givenCallToContextCreateWithoutPowerHintDescThenPowerHintIsNotSetInDriverHandle) {
    ze_context_handle_t hContext;
    ze_context_desc_t ctxtDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_scheduling_hint_exp_desc_t invalidExpContext = {};
    invalidExpContext.stype = ZE_STRUCTURE_TYPE_SCHEDULING_HINT_EXP_DESC;
    ctxtDesc.pNext = &invalidExpContext;
    ze_result_t res = driverHandle->createContext(&ctxtDesc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(0, driverHandle->powerHint);
    L0::Context *context = L0::Context::fromHandle(hContext);
    context->destroy();
}

TEST_F(ContextPowerSavingHintTest, givenOsContextPowerHintMaxAndZePowerSavingHintTypeMaxThenTheyAreEqualAndBothAre100) {
    EXPECT_EQ(NEO::OsContext::getUmdPowerHintMax(), ZE_POWER_SAVING_HINT_TYPE_MAX);
    EXPECT_EQ(NEO::OsContext::getUmdPowerHintMax(), 100u);
}

using ContextTest = Test<DeviceFixture>;

TEST_F(ContextTest, whenCreatingAndDestroyingContextThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = L0::Context::fromHandle(hContext)->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

using ContextMakeMemoryResidentTests = Test<HostPointerManagerFixure>;

TEST_F(ContextMakeMemoryResidentTests,
       givenUknownPointerPassedToMakeMemoryResidentThenInvalidArgumentIsReturned) {
    const size_t size = 4096;
    uint8_t *ptr = new uint8_t[size];
    ze_result_t res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    delete[] ptr;
}

TEST_F(ContextMakeMemoryResidentTests,
       givenValidPointerPassedToMakeMemoryResidentThenSuccessIsReturned) {
    const size_t size = 4096;
    void *ptr = nullptr;
    ze_host_mem_alloc_desc_t hostDesc = {};
    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_result_t res = context->allocSharedMem(device->toHandle(),
                                              &deviceDesc,
                                              &hostDesc,
                                              size,
                                              0,
                                              &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;

    res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    context->freeMem(ptr);
}

TEST_F(ContextMakeMemoryResidentTests,
       givenValidAllocationwithLockedWhenCallingMakeMemoryResidentThenInvalidArgumentIsReturned) {
    const size_t size = 4096;
    void *ptr = nullptr;
    ze_host_mem_alloc_desc_t hostDesc = {};
    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_result_t res = context->allocSharedMem(device->toHandle(),
                                              &deviceDesc,
                                              &hostDesc,
                                              size,
                                              0,
                                              &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    DriverHandle *driverHandle = hostDriverHandle.get();
    auto allocation = driverHandle->getDriverSystemMemoryAllocation(ptr, size, neoDevice->getRootDeviceIndex(), nullptr);
    allocation->setLockedMemory(true);

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;
    res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    context->freeMem(ptr);
}

TEST_F(ContextMakeMemoryResidentTests,
       whenMakingASharedMemoryResidentThenIsAddedToVectorOfResidentAllocations) {
    const size_t size = 4096;
    void *ptr = nullptr;
    ze_host_mem_alloc_desc_t hostDesc = {};
    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_result_t res = context->allocSharedMem(device->toHandle(),
                                              &deviceDesc,
                                              &hostDesc,
                                              size,
                                              0,
                                              &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;
    res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize + 1, currentSize);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t finalSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize, finalSize);

    context->freeMem(ptr);
}

TEST_F(ContextMakeMemoryResidentTests,
       whenMakingADeviceMemoryResidentThenIsNotAddedToVectorOfResidentAllocations) {
    const size_t size = 4096;
    void *ptr = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_result_t res = context->allocDeviceMem(device->toHandle(),
                                              &deviceDesc,
                                              size,
                                              0,
                                              &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;
    res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize, currentSize);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    context->freeMem(ptr);
}

TEST_F(ContextMakeMemoryResidentTests,
       whenMakingASharedMemoryResidentButMemoryInterfaceFailsThenIsNotAddedToVectorOfResidentAllocations) {
    const size_t size = 4096;
    void *ptr = nullptr;
    ze_host_mem_alloc_desc_t hostDesc = {};
    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_result_t res = context->allocSharedMem(device->toHandle(),
                                              &deviceDesc,
                                              &hostDesc,
                                              size,
                                              0,
                                              &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::failed;

    res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_ERROR_DEVICE_LOST, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize, currentSize);

    context->freeMem(ptr);
}

TEST_F(ContextMakeMemoryResidentTests, givenDeviceUnifiedMemoryAndLocalOnlyAllocationModeThenCallMakeMemoryResidentImmediately) {
    DebugManagerStateRestore restorer;
    NEO::debugManager.flags.EnableDeviceUsmAllocationPool.set(0);
    const size_t size = 4096;
    void *ptr = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {};

    auto *driverHandle{hostDriverHandle.get()};
    driverHandle->memoryManager->usmDeviceAllocationMode = NEO::LocalMemAllocationMode::localOnly;
    static_cast<MockMemoryManager *>(driverHandle->memoryManager)->returnFakeAllocation = true;
    hostDriverHandle->svmAllocsManager->cleanupUSMAllocCaches();

    EXPECT_EQ(0U, mockMemoryInterface->makeResidentCalled);
    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;
    ze_result_t res1 = context->allocDeviceMem(device->toHandle(),
                                               &deviceDesc,
                                               size,
                                               0,
                                               &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res1);

    auto allocData{driverHandle->svmAllocsManager->getSVMAlloc(ptr)};
    EXPECT_NE(allocData, nullptr);
    const bool lmemAllocationModeSupported{allocData->gpuAllocations.getDefaultGraphicsAllocation()->storageInfo.localOnlyRequired};
    EXPECT_EQ(mockMemoryInterface->makeResidentCalled, (lmemAllocationModeSupported ? 1U : 0U));
    EXPECT_EQ(ZE_RESULT_SUCCESS, context->freeMem(ptr));

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::outOfMemory;
    ze_result_t res2 = context->allocDeviceMem(device->toHandle(),
                                               &deviceDesc,
                                               size,
                                               0,
                                               &ptr);
    EXPECT_EQ(res2, (lmemAllocationModeSupported ? ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY : ZE_RESULT_SUCCESS));
    EXPECT_EQ(mockMemoryInterface->makeResidentCalled, (lmemAllocationModeSupported ? 2U : 0U));
    EXPECT_EQ(ZE_RESULT_SUCCESS, context->freeMem(ptr));
}

TEST_F(ContextMakeMemoryResidentTests, givenNonDeviceUnifiedMemoryWhenAllocDeviceMemCalledThenMakeMemoryResidentIsNotImmediatelyCalled) {
    DebugManagerStateRestore restorer;
    NEO::debugManager.flags.EnableDeviceUsmAllocationPool.set(0);
    const size_t size = 4096;
    void *ptr = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {};

    auto *driverHandle{hostDriverHandle.get()};
    driverHandle->memoryManager->usmDeviceAllocationMode = NEO::LocalMemAllocationMode::localOnly;
    static_cast<MockMemoryManager *>(driverHandle->memoryManager)->returnFakeAllocation = true;

    auto *origSvmAllocsManager{driverHandle->svmAllocsManager};
    auto fakeAllocationAddr{reinterpret_cast<void *>(0x1234)};

    MockGraphicsAllocation mockUnifiedAllocation{};
    SvmAllocationData allocData(0U);
    allocData.gpuAllocations.addAllocation(&mockUnifiedAllocation);
    allocData.memoryType = InternalMemoryType::notSpecified;

    MockSVMAllocsManager mockSvmAllocsManager{driverHandle->memoryManager};
    mockSvmAllocsManager.createUnifiedMemoryAllocationCallBase = false;
    mockSvmAllocsManager.createUnifiedMemoryAllocationReturnValue = fakeAllocationAddr;
    mockSvmAllocsManager.insertSVMAlloc(fakeAllocationAddr, allocData);
    driverHandle->svmAllocsManager = &mockSvmAllocsManager;

    EXPECT_EQ(0U, mockMemoryInterface->makeResidentCalled);
    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;
    ze_result_t res1 = context->allocDeviceMem(device->toHandle(),
                                               &deviceDesc,
                                               size,
                                               0,
                                               &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res1);
    EXPECT_EQ(mockMemoryInterface->makeResidentCalled, 0U);
    driverHandle->svmAllocsManager = origSvmAllocsManager;
}

struct ContextMakeMemoryResidentAndMigrationTests : public ContextMakeMemoryResidentTests {
    struct MockResidentTestsPageFaultManager : public MockPageFaultManager {
        void moveAllocationToGpuDomain(void *ptr) override {
            moveAllocationToGpuDomainCalledTimes++;
            migratedAddress = ptr;
        }
        void moveAllocationsWithinUMAllocsManagerToGpuDomain(SVMAllocsManager *unifiedMemoryManager) override {
            moveAllocationsWithinUMAllocsManagerToGpuDomainCalled++;
        }
        uint32_t moveAllocationToGpuDomainCalledTimes = 0;
        uint32_t moveAllocationsWithinUMAllocsManagerToGpuDomainCalled = 0;
        void *migratedAddress = nullptr;
    };

    void SetUp() override {
        ContextMakeMemoryResidentTests::SetUp();
        mockMemoryManager = std::make_unique<MockMemoryManager>();
        mockPageFaultManager = new MockResidentTestsPageFaultManager;
        svmManager = std::make_unique<MockSVMAllocsManager>(mockMemoryManager.get());

        mockMemoryManager->pageFaultManager.reset(mockPageFaultManager);
        memoryManager = device->getDriverHandle()->getMemoryManager();
        device->getDriverHandle()->setMemoryManager(mockMemoryManager.get());

        ze_host_mem_alloc_desc_t hostDesc = {};
        ze_device_mem_alloc_desc_t deviceDesc = {};
        ze_result_t res = context->allocSharedMem(device->toHandle(),
                                                  &deviceDesc,
                                                  &hostDesc,
                                                  size,
                                                  0,
                                                  &ptr);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    }

    void TearDown() override {
        device->getDriverHandle()->setMemoryManager(memoryManager);
        ContextMakeMemoryResidentTests::TearDown();
    }

    const size_t size = 4096;
    void *ptr = nullptr;

    std::unique_ptr<MockMemoryManager> mockMemoryManager;
    std::unique_ptr<MockSVMAllocsManager> svmManager;
    MockResidentTestsPageFaultManager *mockPageFaultManager = nullptr;
    NEO::MemoryManager *memoryManager = nullptr;
};

HWTEST_F(ContextMakeMemoryResidentAndMigrationTests,
         whenExecutingCommandListsWithMigrationThenMemoryFromMakeResidentIsMovedToGpu) {
    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;

    ze_result_t res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize + 1, currentSize);

    const ze_command_queue_desc_t desc = {};
    MockCsrHw2<FamilyType> csr(*neoDevice->getExecutionEnvironment(), 0, neoDevice->getDeviceBitfield());
    csr.initializeTagAllocation();
    csr.setupContext(*neoDevice->getDefaultEngine().osContext);

    ze_result_t returnValue;
    L0::CommandQueue *commandQueue = CommandQueue::create(productFamily,
                                                          device,
                                                          &csr,
                                                          &desc,
                                                          true,
                                                          false,
                                                          false,
                                                          returnValue);
    EXPECT_NE(nullptr, commandQueue);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 0u);

    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily,
                                                                     device,
                                                                     NEO::EngineGroupType::copy,
                                                                     0u,
                                                                     returnValue, false));
    auto commandListHandle = commandList->toHandle();
    commandList->close();
    res = commandQueue->executeCommandLists(1, &commandListHandle, nullptr, true, nullptr, nullptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 1u);
    EXPECT_EQ(mockPageFaultManager->migratedAddress, ptr);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    commandQueue->destroy();
    context->freeMem(ptr);
}
HWTEST_F(ContextMakeMemoryResidentAndMigrationTests, whenExecutingKernelWithIndirectAccessThenSharedAllocationsAreMigrated) {

    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;

    ze_result_t res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize + 1, currentSize);

    const ze_command_queue_desc_t desc = {};
    MockCsrHw2<FamilyType> csr(*neoDevice->getExecutionEnvironment(), 0, neoDevice->getDeviceBitfield());
    csr.initializeTagAllocation();
    csr.setupContext(*neoDevice->getDefaultEngine().osContext);

    ze_result_t returnValue;
    L0::CommandQueue *commandQueue = CommandQueue::create(productFamily,
                                                          device,
                                                          &csr,
                                                          &desc,
                                                          true,
                                                          false,
                                                          false,
                                                          returnValue);
    EXPECT_NE(nullptr, commandQueue);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 0u);

    auto commandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<FamilyType::gfxCoreFamily>>>();
    commandList->initialize(device, NEO::EngineGroupType::copy, 0u);
    commandList->unifiedMemoryControls.indirectSharedAllocationsAllowed = true;
    commandList->indirectAllocationsAllowed = true;
    commandList->close();

    auto commandListHandle = commandList->toHandle();
    res = commandQueue->executeCommandLists(1, &commandListHandle, nullptr, true, nullptr, nullptr);

    EXPECT_EQ(mockPageFaultManager->moveAllocationsWithinUMAllocsManagerToGpuDomainCalled, 1u);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    commandQueue->destroy();
    context->freeMem(ptr);
}
HWTEST_F(ContextMakeMemoryResidentAndMigrationTests,
         whenExecutingCommandListsWithNoMigrationThenMemoryFromMakeResidentIsNotMovedToGpu) {
    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;

    ze_result_t res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize + 1, currentSize);

    const ze_command_queue_desc_t desc = {};
    MockCsrHw2<FamilyType> csr(*neoDevice->getExecutionEnvironment(), 0, neoDevice->getDeviceBitfield());
    csr.initializeTagAllocation();
    csr.setupContext(*neoDevice->getDefaultEngine().osContext);

    ze_result_t returnValue;
    L0::CommandQueue *commandQueue = CommandQueue::create(productFamily,
                                                          device,
                                                          &csr,
                                                          &desc,
                                                          true,
                                                          false,
                                                          false,
                                                          returnValue);
    EXPECT_NE(nullptr, commandQueue);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 0u);

    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily,
                                                                     device,
                                                                     NEO::EngineGroupType::copy,
                                                                     0u,
                                                                     returnValue, false));
    auto commandListHandle = commandList->toHandle();
    commandList->close();

    res = commandQueue->executeCommandLists(1, &commandListHandle, nullptr, false, nullptr, nullptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 0u);
    EXPECT_EQ(mockPageFaultManager->migratedAddress, nullptr);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    commandQueue->destroy();
    context->freeMem(ptr);
}

HWTEST_F(ContextMakeMemoryResidentAndMigrationTests,
         whenExecutingImmediateCommandListsHavingSharedAllocationWithMigrationThenMemoryFromMakeResidentIsMovedToGpu) {
    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;

    ze_result_t res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize + 1, currentSize);

    const ze_command_queue_desc_t desc = {};
    MockCsrHw2<FamilyType> csr(*neoDevice->getExecutionEnvironment(), 0, neoDevice->getDeviceBitfield());
    csr.initializeTagAllocation();
    csr.setupContext(*neoDevice->getDefaultEngine().osContext);

    ze_result_t result = ZE_RESULT_SUCCESS;

    std::unique_ptr<L0::CommandList> commandList0(CommandList::createImmediate(productFamily,
                                                                               device,
                                                                               &desc,
                                                                               false,
                                                                               NEO::EngineGroupType::renderCompute,
                                                                               result));
    ASSERT_NE(nullptr, commandList0);

    void *dstBuffer = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_host_mem_alloc_desc_t hostDesc = {};
    result = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, 16384u, 4090u, &dstBuffer);
    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
    CmdListMemoryCopyParams copyParams = {};
    int one = 1;
    result = commandList0->appendMemoryFill(dstBuffer, reinterpret_cast<void *>(&one), sizeof(one), 4090u,
                                            nullptr, 0, nullptr, copyParams);
    EXPECT_EQ(ZE_RESULT_SUCCESS, result);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 1u);
    EXPECT_EQ(mockPageFaultManager->migratedAddress, ptr);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    context->freeMem(ptr);
    context->freeMem(dstBuffer);
}

HWTEST_F(ContextMakeMemoryResidentAndMigrationTests,
         GivenImmediateCommandListWhenExecutingRegularCommandListsHavingSharedAllocationWithMigrationOnImmediateThenMemoryFromMakeResidentIsMovedToGpuOnce) {
    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;

    ze_result_t res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize + 1, currentSize);

    const ze_command_queue_desc_t desc = {};
    MockCsrHw2<FamilyType> csr(*neoDevice->getExecutionEnvironment(), 0, neoDevice->getDeviceBitfield());
    csr.initializeTagAllocation();
    csr.setupContext(*neoDevice->getDefaultEngine().osContext);

    ze_result_t result = ZE_RESULT_SUCCESS;

    std::unique_ptr<L0::CommandList> commandListImmediate(CommandList::createImmediate(productFamily,
                                                                                       device,
                                                                                       &desc,
                                                                                       false,
                                                                                       NEO::EngineGroupType::compute,
                                                                                       result));
    ASSERT_NE(nullptr, commandListImmediate);

    void *dstBuffer = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {};
    ze_host_mem_alloc_desc_t hostDesc = {};
    result = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, 16384u, 4090u, &dstBuffer);
    ASSERT_EQ(ZE_RESULT_SUCCESS, result);

    std::unique_ptr<L0::CommandList> commandListRegular(CommandList::create(productFamily,
                                                                            device,
                                                                            NEO::EngineGroupType::compute,
                                                                            0u,
                                                                            result, false));
    CmdListMemoryCopyParams copyParams = {};
    int one = 1;
    result = commandListRegular->appendMemoryFill(dstBuffer, reinterpret_cast<void *>(&one), sizeof(one), 4090u,
                                                  nullptr, 0, nullptr, copyParams);
    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
    commandListRegular->close();

    auto cmdListHandle = commandListRegular->toHandle();
    result = commandListImmediate->appendCommandLists(1, &cmdListHandle, nullptr, 0, nullptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, result);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 1u);
    EXPECT_EQ(mockPageFaultManager->migratedAddress, ptr);

    mockMemoryInterface->evictResult = NEO::MemoryOperationsStatus::success;
    res = context->evictMemory(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    context->freeMem(ptr);
    context->freeMem(dstBuffer);
}

HWTEST_F(ContextMakeMemoryResidentAndMigrationTests,
         whenExecutingImmediateCommandListsHavingHostAllocationWithMigrationThenMemoryFromMakeResidentIsMovedToGpu) {
    DriverHandle *driverHandle = hostDriverHandle.get();
    size_t previousSize = driverHandle->sharedMakeResidentAllocations.size();

    mockMemoryInterface->makeResidentResult = NEO::MemoryOperationsStatus::success;

    ze_result_t res = context->makeMemoryResident(device, ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t currentSize = driverHandle->sharedMakeResidentAllocations.size();
    EXPECT_EQ(previousSize + 1, currentSize);

    const ze_command_queue_desc_t desc = {};
    MockCsrHw2<FamilyType> csr(*neoDevice->getExecutionEnvironment(), 0, neoDevice->getDeviceBitfield());
    csr.initializeTagAllocation();
    csr.setupContext(*neoDevice->getDefaultEngine().osContext);

    ze_result_t result = ZE_RESULT_SUCCESS;

    std::unique_ptr<L0::CommandList> commandList0(CommandList::createImmediate(productFamily,
                                                                               device,
                                                                               &desc,
                                                                               false,
                                                                               NEO::EngineGroupType::renderCompute,
                                                                               result));
    ASSERT_NE(nullptr, commandList0);

    DebugManagerStateRestore restore;
    debugManager.flags.AllocateSharedAllocationsWithCpuAndGpuStorage.set(true);

    RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
    std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};

    NEO::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::sharedUnifiedMemory, 1, rootDeviceIndices, deviceBitfields);
    auto sharedPtr = svmManager->createSharedUnifiedMemoryAllocation(4096u, unifiedMemoryProperties, device);
    EXPECT_NE(nullptr, sharedPtr);

    auto allocation = svmManager->getSVMAlloc(sharedPtr);
    auto gpuAllocation = allocation->gpuAllocations.getGraphicsAllocation(mockRootDeviceIndex);

    auto &commandContainer = commandList0->getCmdContainer();
    commandContainer.addToResidencyContainer(gpuAllocation);
    commandContainer.addToResidencyContainer(allocation->cpuAllocation);

    void *dstBuffer = nullptr;
    ze_host_mem_alloc_desc_t hostDesc = {};
    result = context->allocHostMem(&hostDesc, 4096u, 0u, &dstBuffer);
    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
    CmdListMemoryCopyParams copyParams = {};
    int one = 1;
    result = commandList0->appendMemoryFill(dstBuffer, reinterpret_cast<void *>(&one), sizeof(one), 4090u,
                                            nullptr, 0, nullptr, copyParams);
    EXPECT_EQ(ZE_RESULT_SUCCESS, result);

    EXPECT_EQ(mockPageFaultManager->moveAllocationToGpuDomainCalledTimes, 3u);

    context->freeMem(ptr);
    svmManager->freeSVMAlloc(sharedPtr);
    context->freeMem(dstBuffer);
}

TEST_F(ContextTest, whenGettingDriverThenDriverIsRetrievedSuccessfully) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    L0::DriverHandle *driverHandleFromContext = contextImp->getDriverHandle();
    EXPECT_EQ(driverHandleFromContext, driverHandle.get());

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

HWTEST2_F(ContextTest, whenCreatingContextWithSvmHeapReservationTrueThenContextSettingsSetCorrectly, IsNotMTL) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    EXPECT_TRUE(contextImp->settings.enableSvmHeapReservation);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemInterfacesThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingQueryVirtualMemPageSizeCorrectAlignmentIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    size_t size = 1024;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(pagesize, MemoryConstants::pageSize64k);

    size = MemoryConstants::pageSize2M - 1000;
    pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(pagesize, MemoryConstants::pageSize64k);

    size = MemoryConstants::pageSize2M + 1000;
    pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(pagesize, MemoryConstants::pageSize64k);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingQueryVirtualMemPageSizeWithStartAddressThenCorrectAlignmentIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    size_t size = 1024;
    size_t pagesize = 0u;

    const auto &gfxPartition = driverHandle->getMemoryManager()->getGfxPartition(neoDevice->getRootDeviceIndex());
    for (uint32_t heapIndex = static_cast<uint32_t>(HeapIndex::heapInternalDeviceMemory); heapIndex < static_cast<uint32_t>(HeapIndex::totalHeaps); ++heapIndex) {
        if (gfxPartition->isHeapInitialized(static_cast<HeapIndex>(heapIndex)) == false) {
            continue;
        }
        auto heapBase = gfxPartition->getHeapBase(static_cast<HeapIndex>(heapIndex));
        auto heapAlignment = gfxPartition->getHeapAllocationAlignment(static_cast<HeapIndex>(heapIndex));

        uint32_t offset = 0;
        if (HeapIndex::heapSvm == static_cast<HeapIndex>(heapIndex)) {
            offset = 32; // Adding offset to avoid having heapSvm base to be at 0 and conflicting with pStart = nullptr case
        }
        void *pStart = reinterpret_cast<void *>(heapBase + offset);

        res = contextImp->queryVirtualMemPageSizeWithStartAddress(device, pStart, size, &pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        EXPECT_EQ(pagesize, heapAlignment);
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingQueryVirtualMemPageSizeWithInvalidStartAddressThenDefaultAlignmentIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    struct MockContextImp : public ContextImp {
        using ContextImp::getPageAlignedSizeRequired;
    };

    MockContextImp *contextImp = static_cast<MockContextImp *>(L0::Context::fromHandle(hContext));

    size_t size = 1024;
    size_t defaultPageSize = 0u;

    auto alignedSize = contextImp->getPageAlignedSizeRequired(nullptr, size, nullptr, &defaultPageSize);
    EXPECT_EQ(defaultPageSize, MemoryConstants::pageSize64k);

    const auto &gfxPartition = driverHandle->getMemoryManager()->getGfxPartition(neoDevice->getRootDeviceIndex());
    uint64_t maxHeapLimit = 0;
    for (uint32_t heapIndex = static_cast<uint32_t>(HeapIndex::heapInternalDeviceMemory); heapIndex < static_cast<uint32_t>(HeapIndex::totalHeaps); ++heapIndex) {
        maxHeapLimit = std::max(maxHeapLimit, gfxPartition->getHeapLimit(static_cast<HeapIndex>(heapIndex)));
    }
    void *pStart = reinterpret_cast<void *>(maxHeapLimit + 64);

    size_t pageSize = 0u;
    contextImp->getPageAlignedSizeRequired(pStart, alignedSize, nullptr, &pageSize);
    EXPECT_EQ(pageSize, defaultPageSize);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingPhysicalMemInterfacesThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    size_t size = 1024;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingPhysicalMemCreateWithInvalisSizeThenUnsupportedSizeReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    size_t size = 1024;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_DEVICE, 10};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);

    descMem.flags = ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_HOST;
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingDestroyPhysicalMemWithIncorrectPointerThenMemoryNotFreed) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    size_t size = 1024;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getPhysicalMemoryAllocationMap().size()), 0);

    res = contextImp->destroyPhysicalMem(nullptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getPhysicalMemoryAllocationMap().size()), 0);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(static_cast<int>(driverHandle->getMemoryManager()->getPhysicalMemoryAllocationMap().size()), 0);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingMappingVirtualInterfacesOnPhysicalDeviceMemoryThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_DEVICE, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    std::vector<ze_memory_access_attribute_t> memoryAccessFlags = {
        ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE, ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY,
        ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};

    for (auto accessFlags : memoryAccessFlags) {
        res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->unMapVirtualMem(ptr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    }

    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->setVirtualMemAccessAttribute(ptr, pagesize, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t outAccess = {};
    size_t outSize = 0;
    res = contextImp->getVirtualMemAccessAttribute(ptr, pagesize, &outAccess, &outSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(pagesize, outSize);

    res = contextImp->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingMappingVirtualInterfacesOnPhysicalDeviceMemoryThenMakeResidentIsCalledWithForcePagingFenceTrue) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    NEO::MockMemoryOperations *mockMemoryInterface = static_cast<NEO::MockMemoryOperations *>(
        device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface.get());

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_DEVICE, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size_t offset = 0;

    std::vector<ze_memory_access_attribute_t> memoryAccessFlags = {
        ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE, ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY,
        ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};

    for (auto accessFlags : memoryAccessFlags) {
        EXPECT_FALSE(mockMemoryInterface->makeResidentForcePagingFenceValue);
        res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, accessFlags);
        EXPECT_TRUE(mockMemoryInterface->makeResidentForcePagingFenceValue);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);

        res = contextImp->unMapVirtualMem(ptr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);

        mockMemoryInterface->makeResidentForcePagingFenceValue = false;
    }

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingMappingVirtualInterfacesOnPhysicalHostMemoryThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_HOST, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    std::vector<ze_memory_access_attribute_t> memoryAccessFlags = {
        ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE, ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY,
        ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};

    for (auto accessFlags : memoryAccessFlags) {
        res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->unMapVirtualMem(ptr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    }

    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->setVirtualMemAccessAttribute(ptr, pagesize, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t outAccess = {};
    size_t outSize = 0;
    res = contextImp->getVirtualMemAccessAttribute(ptr, pagesize, &outAccess, &outSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(pagesize, outSize);

    res = contextImp->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemorySetAttributeWithValidEnumerationsThenSuccessisReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    std::vector<ze_memory_access_attribute_t> memoryAccessFlags = {
        ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE, ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY,
        ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};
    ze_memory_access_attribute_t previousAccess = {};
    size_t outSize = 0;
    for (auto accessFlags : memoryAccessFlags) {
        previousAccess = access;
        res = contextImp->setVirtualMemAccessAttribute(ptr, pagesize, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->getVirtualMemAccessAttribute(ptr, pagesize, &access, &outSize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        EXPECT_EQ(outSize, pagesize);
        EXPECT_NE(previousAccess, access);
        EXPECT_EQ(accessFlags, access);
    }

    res = contextImp->getVirtualMemAccessAttribute(ptr, pagesize, &access, &outSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemorySetAttributeWithInvalidValuesThenFailureReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_FORCE_UINT32};
    res = contextImp->setVirtualMemAccessAttribute(ptr, pagesize, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ENUMERATION, res);

    res = contextImp->setVirtualMemAccessAttribute(pStart, pagesize, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryGetAttributeWithInvalidValuesThenFailureReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    size_t outSize = 0;
    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};
    res = contextImp->getVirtualMemAccessAttribute(pStart, pagesize, &access, &outSize);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryFreeWithInvalidValuesThenFailuresReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    size_t invalidSize = 10u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    res = contextImp->freeVirtualMem(ptr, invalidSize);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->freeVirtualMem(pStart, pagesize);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
    pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);

    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    res = contextImp->freeVirtualMem(ptr, invalidSize);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

class ReserveMemoryManagerMock : public NEO::MemoryManager {
  public:
    ReserveMemoryManagerMock(NEO::ExecutionEnvironment &executionEnvironment) : NEO::MemoryManager(executionEnvironment) {}
    NEO::GraphicsAllocation *createGraphicsAllocationFromMultipleSharedHandles(const std::vector<osHandle> &handles, AllocationProperties &properties, bool requireSpecificBitness, bool isHostIpcAllocation, bool reuseSharedAllocation, void *mapPointer) override { return nullptr; }
    NEO::GraphicsAllocation *createGraphicsAllocationFromSharedHandle(const OsHandleData &osHandleData, const AllocationProperties &properties, bool requireSpecificBitness, bool isHostIpcAllocation, bool reuseSharedAllocation, void *mapPointer) override { return nullptr; }
    void addAllocationToHostPtrManager(NEO::GraphicsAllocation *memory) override{};
    void removeAllocationFromHostPtrManager(NEO::GraphicsAllocation *memory) override{};
    AllocationStatus populateOsHandles(NEO::OsHandleStorage &handleStorage, uint32_t rootDeviceIndex) override { return AllocationStatus::Success; };
    void cleanOsHandles(NEO::OsHandleStorage &handleStorage, uint32_t rootDeviceIndex) override{};
    void freeGraphicsMemoryImpl(NEO::GraphicsAllocation *gfxAllocation) override{};
    void freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation, bool isImportedAllocation) override{};
    uint64_t getSystemSharedMemory(uint32_t rootDeviceIndex) override { return 0; };
    uint64_t getLocalMemorySize(uint32_t rootDeviceIndex, uint32_t deviceBitfield) override { return 0; };
    double getPercentOfGlobalMemoryAvailable(uint32_t rootDeviceIndex) override { return 0; }
    AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, const RootDeviceIndicesContainer &rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override {
        if (failReserveGpuAddress) {
            return {};
        }
        return AddressRange{requiredStartAddress, size};
    }
    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, const RootDeviceIndicesContainer &rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
        if (failReserveGpuAddress) {
            return {};
        }
        return AddressRange{requiredStartAddress, size};
    }
    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
        return selectAlignmentAndHeap(0ull, size, heap);
    }
    size_t selectAlignmentAndHeap(const uint64_t pStart, size_t size, HeapIndex *heap) override {

        // Always default to HEAP STANDARD and page size 64Kb.
        *heap = HeapIndex::heapStandard;
        size_t pageSizeAlignment = MemoryConstants::pageSize64k;

        // If the user provides a start address, we try to find the heap and page size alignment based on that address.
        if (pStart != 0ull) {
            auto rootDeviceIndex = 0u;
            auto gfxPartition = getGfxPartition(rootDeviceIndex);
            if (gfxPartition->getHeapIndexAndPageSizeBasedOnAddress(pStart, *heap, pageSizeAlignment)) {
                return pageSizeAlignment;
            }
        }
        return pageSizeAlignment;
    }
    void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override {
        if (failReserveCpuAddress) {
            return {};
        }
        return AddressRange{requiredStartAddress, size};
    }
    void freeCpuAddress(AddressRange addressRange) override{};
    NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
    NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
    NEO::GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const NEO::AllocationData &allocationData) override { return nullptr; };
    NEO::GraphicsAllocation *allocateUSMHostGraphicsMemory(const NEO::AllocationData &allocationData) override { return nullptr; };
    NEO::GraphicsAllocation *allocateGraphicsMemory64kb(const NEO::AllocationData &allocationData) override { return nullptr; };
    NEO::GraphicsAllocation *allocate32BitGraphicsMemoryImpl(const NEO::AllocationData &allocationData) override { return nullptr; };
    NEO::GraphicsAllocation *allocateGraphicsMemoryInDevicePool(const NEO::AllocationData &allocationData, AllocationStatus &status) override { return nullptr; };
    NEO::GraphicsAllocation *allocateGraphicsMemoryWithGpuVa(const NEO::AllocationData &allocationData) override { return nullptr; };
    GraphicsAllocation *allocatePhysicalDeviceMemory(const AllocationData &allocationData, AllocationStatus &status) override { return nullptr; };
    GraphicsAllocation *allocatePhysicalLocalDeviceMemory(const AllocationData &allocationData, AllocationStatus &status) override { return nullptr; };
    GraphicsAllocation *allocatePhysicalHostMemory(const AllocationData &allocationData, AllocationStatus &status) override { return nullptr; };
    bool unMapPhysicalDeviceMemoryFromVirtualMemory(GraphicsAllocation *physicalAllocation, uint64_t gpuRange, size_t bufferSize, OsContext *osContext, uint32_t rootDeviceIndex) override { return false; };
    bool unMapPhysicalHostMemoryFromVirtualMemory(MultiGraphicsAllocation &multiGraphicsAllocation, GraphicsAllocation *physicalAllocation, uint64_t gpuRange, size_t bufferSize) override { return false; };
    bool mapPhysicalDeviceMemoryToVirtualMemory(GraphicsAllocation *physicalAllocation, uint64_t gpuRange, size_t bufferSize, const MemoryFlags *memoryflags) override {
        if (failMapVirtualMemory) {
            return false;
        } else {
            return true;
        }
    };
    bool mapPhysicalHostMemoryToVirtualMemory(RootDeviceIndicesContainer &rootDeviceIndices, MultiGraphicsAllocation &multiGraphicsAllocation, GraphicsAllocation *physicalAllocation, uint64_t gpuRange, size_t bufferSize) override {
        if (failMapVirtualMemory) {
            return false;
        } else {
            return true;
        }
    };

    NEO::GraphicsAllocation *allocatePhysicalGraphicsMemory(const AllocationProperties &properties) override {

        if (failAllocatePhysicalGraphicsMemory) {
            return nullptr;
        }
        mockAllocation.reset(new NEO::MockGraphicsAllocation(const_cast<void *>(buffer), size));
        return mockAllocation.get();
    }

    NEO::GraphicsAllocation *allocateGraphicsMemoryForImageImpl(const NEO::AllocationData &allocationData, std::unique_ptr<Gmm> gmm) override { return nullptr; };
    NEO::GraphicsAllocation *allocateMemoryByKMD(const NEO::AllocationData &allocationData) override { return nullptr; };
    void *lockResourceImpl(NEO::GraphicsAllocation &graphicsAllocation) override { return nullptr; };
    void unlockResourceImpl(NEO::GraphicsAllocation &graphicsAllocation) override{};
    bool peek32bit() override {
        return this->is32bit;
    }

    bool failReserveGpuAddress = true;
    bool failReserveCpuAddress = true;
    bool failMapVirtualMemory = true;
    bool failAllocatePhysicalGraphicsMemory = true;
    bool is32bit = false;
    void *buffer = nullptr;
    size_t size = 0;
    std::unique_ptr<NEO::GraphicsAllocation> mockAllocation;
};

HWTEST2_F(ContextTest, whenCallingVirtualMemReserveWithPStartInSvmRangeWithSuccessfulAllocationThenSuccessReturned, IsNotMTL) {
    ze_context_handle_t hContext{};
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = reinterpret_cast<void *>(0x1234);
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    auto reserveMemoryManager = std::make_unique<ReserveMemoryManagerMock>(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    reserveMemoryManager->failReserveCpuAddress = false;
    driverHandle->setMemoryManager(reserveMemoryManager.get());
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(reserveMemoryManager->getVirtualMemoryReservationMap().size(), 0u);
    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    driverHandle->setMemoryManager(memoryManager);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemReserveWithPStartAboveSvmRangeWithSuccessfulAllocationThenSuccessReturned) {
    ze_context_handle_t hContext{};
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    auto reserveMemoryManager = std::make_unique<ReserveMemoryManagerMock>(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    reserveMemoryManager->failReserveGpuAddress = false;
    driverHandle->setMemoryManager(reserveMemoryManager.get());

    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
    void *pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);
    size_t size = 4096u;
    void *ptr = nullptr;

    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(reserveMemoryManager->getVirtualMemoryReservationMap().size(), 0u);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    driverHandle->setMemoryManager(memoryManager);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenUsingOffsetsIntoReservedVirtualMemoryThenMappingIsSuccessful) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize * 2, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    std::vector<ze_memory_access_attribute_t> memoryAccessFlags = {
        ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE, ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY,
        ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};

    void *offsetAddr =
        reinterpret_cast<void *>(reinterpret_cast<uint64_t>(ptr) + pagesize);

    for (auto accessFlags : memoryAccessFlags) {
        res = contextImp->mapVirtualMem(offsetAddr, pagesize, mem, offset, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->setVirtualMemAccessAttribute(offsetAddr, pagesize, access);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);

        ze_memory_access_attribute_t outAccess = {};
        size_t outSize = 0;
        res = contextImp->getVirtualMemAccessAttribute(offsetAddr, pagesize, &outAccess, &outSize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        EXPECT_EQ(pagesize * 2, outSize);

        res = contextImp->unMapVirtualMem(offsetAddr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    }

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize * 2);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenUsingOffsetsIntoReservedVirtualMemoryWithMultiplePhysicalMemoryThenMappingIsSuccessful) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize * 2, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    ze_physical_mem_handle_t secondHalfMem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &secondHalfMem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    std::vector<ze_memory_access_attribute_t> memoryAccessFlags = {
        ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE, ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY,
        ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};

    void *offsetAddr =
        reinterpret_cast<void *>(reinterpret_cast<uint64_t>(ptr) + pagesize);

    for (auto accessFlags : memoryAccessFlags) {
        res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->mapVirtualMem(offsetAddr, pagesize, secondHalfMem, offset, accessFlags);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->setVirtualMemAccessAttribute(ptr, pagesize, access);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->setVirtualMemAccessAttribute(offsetAddr, pagesize, access);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);

        ze_memory_access_attribute_t outAccess = {};
        size_t outSize = 0;
        res = contextImp->getVirtualMemAccessAttribute(ptr, pagesize, &outAccess, &outSize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        EXPECT_EQ(pagesize * 2, outSize);
        res = contextImp->getVirtualMemAccessAttribute(offsetAddr, pagesize, &outAccess, &outSize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        EXPECT_EQ(pagesize * 2, outSize);

        res = contextImp->unMapVirtualMem(ptr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
        res = contextImp->unMapVirtualMem(offsetAddr, pagesize);
        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    }

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroyPhysicalMem(secondHalfMem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize * 2);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

HWTEST2_F(ContextTest, whenCallingVirtualMemoryReservationWhenOutOfMemoryThenOutOfMemoryReturned, IsNotMTL) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 0u;
    void *ptr = nullptr;
    size_t pageSize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    NEO::MemoryManager *failingReserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(failingReserveMemoryManager);
    res = contextImp->reserveVirtualMem(pStart, pageSize, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);

    pStart = reinterpret_cast<void *>(0x1234);
    res = contextImp->reserveVirtualMem(pStart, pageSize, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);

    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
    pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);
    res = contextImp->reserveVirtualMem(pStart, pageSize, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);

    driverHandle->setMemoryManager(memoryManager);
    delete failingReserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithInvalidArgumentsThenUnsupportedSizeReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 64u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);

    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    NEO::MemoryManager *failingReserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(failingReserveMemoryManager);
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);

    driverHandle->setMemoryManager(memoryManager);
    delete failingReserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithInvalidMultiPageSizeInArgumentsThenUnsupportedSizeReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 64u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);

    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    NEO::MemoryManager *failingReserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(failingReserveMemoryManager);

    size = pagesize * 3 + 10;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);

    size = pagesize * 2 + 1;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);

    driverHandle->setMemoryManager(memoryManager);
    delete failingReserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationOnNonSvmHeapWithUnAlignedSizeInArgumentsThenUnsupportedSizeReturned) {

    DebugManagerStateRestore restore;
    NEO::debugManager.flags.EnableReservingInSvmRange.set(0);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 64u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);

    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    NEO::MemoryManager *failingReserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(failingReserveMemoryManager);

    size = pagesize + MemoryConstants::pageSize;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);

    driverHandle->setMemoryManager(memoryManager);
    delete failingReserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationOnNonSvmHeapUsingPstartWithUnAlignedSizeInArgumentsThenSuccessReturned) {

    DebugManagerStateRestore restore;
    NEO::debugManager.flags.EnableReservingInSvmRange.set(0);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    auto heapIndex = HeapIndex::heapStandard64KB;
    const auto &gfxPartition = driverHandle->getMemoryManager()->getGfxPartition(neoDevice->getRootDeviceIndex());

    if (gfxPartition->isHeapInitialized(heapIndex) == false) {
        GTEST_SKIP();
    }

    auto heapBase = gfxPartition->getHeapBase(static_cast<HeapIndex>(heapIndex));
    auto heapAlignment = gfxPartition->getHeapAllocationAlignment(static_cast<HeapIndex>(heapIndex));

    void *pStart = addrToPtr(heapBase);
    size_t size = 64u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSizeWithStartAddress(device, pStart, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(pagesize, heapAlignment);

    size = pagesize + MemoryConstants::pageSize;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationOnNonSvmHeapUsingPstartWithAlignedSizeInArgumentsThenSuccessReturned) {

    DebugManagerStateRestore restore;
    NEO::debugManager.flags.EnableReservingInSvmRange.set(0);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    auto heapIndex = HeapIndex::heapStandard;
    const auto &gfxPartition = driverHandle->getMemoryManager()->getGfxPartition(neoDevice->getRootDeviceIndex());

    auto heapBase = gfxPartition->getHeapBase(static_cast<HeapIndex>(heapIndex));
    auto heapAlignment = gfxPartition->getHeapAllocationAlignment(static_cast<HeapIndex>(heapIndex));

    void *pStart = addrToPtr(heapBase);
    size_t size = 64u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSizeWithStartAddress(device, pStart, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(pagesize, heapAlignment);

    size = pagesize;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationOnNonSvmHeapWithAlignedSizeInArgumentsThenSuccessReturned) {

    DebugManagerStateRestore restore;
    NEO::debugManager.flags.EnableReservingInSvmRange.set(0);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 64u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    size = pagesize;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithValidMultiPageSizeInArgumentsThenSuccessReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 64u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);

    size = pagesize * 3;

    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithOverlappingReservationRangeThenSuccessReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 0;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    size = 16 * pagesize;

    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    void *newPStart = addrToPtr(castToUint64(ptr) + pagesize);
    void *newPtr = nullptr;
    res = contextImp->reserveVirtualMem(newPStart, pagesize, &newPtr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, newPtr);

    EXPECT_NE(ptr, newPtr);

    res = contextImp->freeVirtualMem(newPtr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

class MockCpuInfoOverrideVirtualAddressSize {
  public:
    class MockCpuInfo : public CpuInfo {
      public:
        using CpuInfo::cpuFlags;
        using CpuInfo::virtualAddressSize;
    } *mockCpuInfo = reinterpret_cast<MockCpuInfo *>(const_cast<CpuInfo *>(&CpuInfo::getInstance()));

    MockCpuInfoOverrideVirtualAddressSize(uint32_t newCpuVirtualAddressSize) {
        virtualAddressSizeSave = mockCpuInfo->getVirtualAddressSize();
        mockCpuInfo->virtualAddressSize = newCpuVirtualAddressSize;
    }

    ~MockCpuInfoOverrideVirtualAddressSize() {
        mockCpuInfo->virtualAddressSize = virtualAddressSizeSave;
    }

    uint32_t virtualAddressSizeSave = 0;
};

HWTEST2_F(ContextTest, Given32BitCpuAddressWidthWhenCallingVirtualMemoryReservationCorrectAllocationMethodIsSelected, IsNotMTL) {
    MockCpuInfoOverrideVirtualAddressSize overrideCpuInfo(32);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = addrToPtr(0x1234);
    size_t size = MemoryConstants::pageSize2M;
    void *ptr = nullptr;

    auto reserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    reserveMemoryManager->is32bit = true;
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(reserveMemoryManager);

    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
    reserveMemoryManager->failReserveCpuAddress = false;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    pStart = addrToPtr(maxNBitValue(32) + 0x1234);

    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
    reserveMemoryManager->failReserveGpuAddress = false;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->setMemoryManager(memoryManager);
    delete reserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

HWTEST2_F(ContextTest, Given48BitCpuAddressWidthWhenCallingVirtualMemoryReservationCorrectAllocationMethodIsSelected, IsNotMTL) {
    MockCpuInfoOverrideVirtualAddressSize overrideCpuInfo(48);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = addrToPtr(0x1234);
    size_t size = MemoryConstants::pageSize2M;
    void *ptr = nullptr;

    auto reserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(reserveMemoryManager);

    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
    reserveMemoryManager->failReserveCpuAddress = false;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    pStart = addrToPtr(maxNBitValue(47) + 0x1234);

    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
    reserveMemoryManager->failReserveGpuAddress = false;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->setMemoryManager(memoryManager);
    delete reserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

HWTEST2_F(ContextTest, Given57BitCpuAddressWidthWhenCallingVirtualMemoryReservationCorrectAllocationMethodIsSelected, IsNotMTL) {

    MockCpuInfoOverrideVirtualAddressSize overrideCpuInfo(57);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = addrToPtr(0x1234);
    size_t size = MemoryConstants::pageSize2M;
    void *ptr = nullptr;

    auto reserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(reserveMemoryManager);

    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
    reserveMemoryManager->failReserveCpuAddress = false;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    pStart = addrToPtr(maxNBitValue(56) + 0x1234);
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);

    reserveMemoryManager->failReserveGpuAddress = false;
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(ptr, nullptr);
    EXPECT_NE(ptr, pStart);

    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->setMemoryManager(memoryManager);
    delete reserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithUnAlignedPstartThenNearbyAlignedPstartIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    auto reserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(reserveMemoryManager);
    reserveMemoryManager->failReserveGpuAddress = false;

    size_t size = MemoryConstants::pageSize;
    void *ptr = nullptr;

    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
    void *pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);

    // pStart is not aligned to any pagesize. The reserveVirtualMem will properly align it.
    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(ptr, nullptr);
    EXPECT_NE(ptr, pStart);
    res = contextImp->freeVirtualMem(ptr, size);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->setMemoryManager(memoryManager);
    delete reserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingPhysicalMemoryAllocateWhenOutOfMemoryThenOutofMemoryReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    size_t size = 0u;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    NEO::MemoryManager *failingReserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(failingReserveMemoryManager);
    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
    driverHandle->setMemoryManager(memoryManager);
    delete failingReserveMemoryManager;

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingMapVirtualDeviceMemoryWithFailedMapThenOutOfMemoryreturned) {
    ze_context_handle_t hContext;
    std::unique_ptr<MockMemoryManager> mockMemoryManager;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    mockMemoryManager = std::make_unique<MockMemoryManager>();

    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(mockMemoryManager.get());

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_DEVICE, pagesize};
    ze_physical_mem_handle_t mem = {};

    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;
    mockMemoryManager->failMapPhysicalToVirtualMemory = true;
    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->setMemoryManager(memoryManager);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingMapVirtualHostMemoryWithFailedMapThenOutOfMemoryreturned) {
    ze_context_handle_t hContext;
    std::unique_ptr<MockMemoryManager> mockMemoryManager;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    mockMemoryManager = std::make_unique<MockMemoryManager>();

    auto memoryManager = driverHandle->getMemoryManager();
    driverHandle->setMemoryManager(mockMemoryManager.get());

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_HOST, pagesize};
    ze_physical_mem_handle_t mem = {};

    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;
    mockMemoryManager->failMapPhysicalToVirtualMemory = true;
    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    driverHandle->setMemoryManager(memoryManager);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingMapVirtualMemoryWithInvalidValuesThenFailureReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize * 2, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    ze_physical_mem_handle_t invalidMem = {};

    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    res = contextImp->mapVirtualMem(ptr, pagesize, invalidMem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->mapVirtualMem(nullptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = contextImp->mapVirtualMem(ptr, pagesize * 4, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    void *offsetAddr =
        reinterpret_cast<void *>(reinterpret_cast<uint64_t>(ptr) + pagesize);
    res = contextImp->mapVirtualMem(offsetAddr, pagesize * 2, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    // Test coverage for case when "if (virtualMemoryReservation)" is false.
    res = contextImp->unMapVirtualMem(nullptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->mapVirtualMem(ptr, 0u, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_ALIGNMENT, res);

    access = ZE_MEMORY_ACCESS_ATTRIBUTE_FORCE_UINT32;
    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ENUMERATION, res);

    access = ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE;
    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    // Test unsupported page size alignment: not covered in API validation layer.
    res = contextImp->unMapVirtualMem(ptr, pagesize + pagesize / 2);
    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_ALIGNMENT, res);

    res = contextImp->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize * 2);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingUnmapVirtualMemoryWithFailedUnmapThenUnknownErrorReturned) {
    ze_context_handle_t hContext;
    std::unique_ptr<MockMemoryManager> mockMemManager;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface =
        std::make_unique<NEO::MockMemoryOperations>();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    // Use a mock memory manager that returns a failure for unmapVirtualMem.
    auto memManager = driverHandle->getMemoryManager();

    mockMemManager = std::make_unique<MockMemoryManager>();
    driverHandle->setMemoryManager(mockMemManager.get());

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;

    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, ZE_PHYSICAL_MEM_FLAG_ALLOCATE_ON_DEVICE, pagesize};
    ze_physical_mem_handle_t mem = {};

    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;

    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Set the mock memory manager to return an error for unmapVirtualMem.
    mockMemManager->failUnMapPhysicalToVirtualMemory = true;
    res = contextImp->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_ERROR_UNKNOWN, res);

    // Reset the mock memory manager to not fail unmapVirtualMem.
    mockMemManager->failUnMapPhysicalToVirtualMemory = false;

    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Reset the memory manager to the original one.
    driverHandle->setMemoryManager(memManager);

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

using IsAtMostProductDG1 = IsAtMostProduct<IGFX_DG1>;

HWTEST2_F(ContextTest, WhenCreatingImageThenSuccessIsReturned, IsAtMostProductDG1) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    ze_image_handle_t image = {};
    ze_image_desc_t imageDesc = {};
    imageDesc.stype = ZE_STRUCTURE_TYPE_IMAGE_DESC;

    res = contextImp->createImage(device, &imageDesc, &image);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, image);

    Image::fromHandle(image)->destroy();

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

HWTEST_F(ContextTest, givenBindlessModeDisabledWhenMakeImageResidentAndEvictThenImageImplicitArgsAllocationIsNotMadeResidentAndEvicted) {
    if (!device->getNEODevice()->getRootDeviceEnvironment().getReleaseHelper() ||
        !device->getNEODevice()->getDeviceInfo().imageSupport) {
        GTEST_SKIP();
    }

    DebugManagerStateRestore restore;
    NEO::debugManager.flags.UseBindlessMode.set(0);

    auto rootDeviceIndex = device->getNEODevice()->getRootDeviceIndex();
    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[rootDeviceIndex]->bindlessHeapsHelper.reset();

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    auto mockMemoryOperationsInterface = new NEO::MockMemoryOperations();
    mockMemoryOperationsInterface->captureGfxAllocationsForMakeResident = true;
    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface.reset(mockMemoryOperationsInterface);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    ze_image_handle_t image = {};
    ze_image_desc_t imageDesc = {};
    imageDesc.stype = ZE_STRUCTURE_TYPE_IMAGE_DESC;

    res = contextImp->createImage(device, &imageDesc, &image);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, image);

    {
        contextImp->makeImageResident(device, image);
        EXPECT_EQ(1, mockMemoryOperationsInterface->makeResidentCalledCount);
        auto allocIter = std::find(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.begin(),
                                   mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(),
                                   Image::fromHandle(image)->getImplicitArgsAllocation());
        EXPECT_EQ(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(), allocIter);
    }

    {
        contextImp->evictImage(device, image);
        EXPECT_EQ(1, mockMemoryOperationsInterface->evictCalledCount);
        auto allocIter = std::find(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.begin(),
                                   mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(),
                                   Image::fromHandle(image)->getImplicitArgsAllocation());
        EXPECT_EQ(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(), allocIter);
    }

    Image::fromHandle(image)->destroy();

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

HWTEST_F(ContextTest, givenBindlessImageWhenMakeImageResidentAndEvictThenImageImplicitArgsAllocationIsMadeResidentAndEvicted) {
    if (!device->getNEODevice()->getRootDeviceEnvironment().getReleaseHelper() ||
        !device->getNEODevice()->getDeviceInfo().imageSupport) {
        GTEST_SKIP();
    }

    DebugManagerStateRestore restore;
    NEO::debugManager.flags.UseBindlessMode.set(1);

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    auto mockMemoryOperationsInterface = new NEO::MockMemoryOperations();
    mockMemoryOperationsInterface->captureGfxAllocationsForMakeResident = true;
    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface.reset(mockMemoryOperationsInterface);
    auto bindlessHelper = new MockBindlesHeapsHelper(device->getNEODevice(),
                                                     device->getNEODevice()->getNumGenericSubDevices() > 1);
    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[device->getNEODevice()->getRootDeviceIndex()]->bindlessHeapsHelper.reset(bindlessHelper);
    // Reset ResidentCalledCount to 0 since bindlessHeapsHelper constructor changes the value.
    mockMemoryOperationsInterface->makeResidentCalledCount = 0;

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    ze_image_handle_t image = {};
    ze_image_bindless_exp_desc_t bindlessExtDesc = {};
    bindlessExtDesc.stype = ZE_STRUCTURE_TYPE_BINDLESS_IMAGE_EXP_DESC;
    bindlessExtDesc.pNext = nullptr;
    bindlessExtDesc.flags = ZE_IMAGE_BINDLESS_EXP_FLAG_BINDLESS;

    ze_image_desc_t imageDesc = {};
    imageDesc.stype = ZE_STRUCTURE_TYPE_IMAGE_DESC;
    imageDesc.pNext = &bindlessExtDesc;

    res = contextImp->createImage(device, &imageDesc, &image);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, image);

    {
        contextImp->makeImageResident(device, image);
        EXPECT_EQ(2, mockMemoryOperationsInterface->makeResidentCalledCount);
        auto allocIter = std::find(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.begin(),
                                   mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(),
                                   Image::fromHandle(image)->getImplicitArgsAllocation());
        EXPECT_NE(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(), allocIter);
    }

    {
        contextImp->evictImage(device, image);
        EXPECT_EQ(2, mockMemoryOperationsInterface->evictCalledCount);
        auto allocIter = std::find(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.begin(),
                                   mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(),
                                   Image::fromHandle(image)->getImplicitArgsAllocation());
        EXPECT_EQ(mockMemoryOperationsInterface->gfxAllocationsForMakeResident.end(), allocIter);
    }

    Image::fromHandle(image)->destroy();

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, givenImageWhenEvictImageWithoutMakeResidentThenSuccessReturned) {
    if (!device->getNEODevice()->getRootDeviceEnvironment().getReleaseHelper() ||
        !device->getNEODevice()->getDeviceInfo().imageSupport) {
        GTEST_SKIP();
    }

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    auto mockMemoryOperationsInterface = new NEO::MockMemoryOperations();
    mockMemoryOperationsInterface->captureGfxAllocationsForMakeResident = true;
    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface.reset(mockMemoryOperationsInterface);
    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    ze_image_handle_t image = {};

    ze_image_desc_t imageDesc = {};
    imageDesc.stype = ZE_STRUCTURE_TYPE_IMAGE_DESC;

    res = contextImp->createImage(device, &imageDesc, &image);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, image);

    EXPECT_EQ(contextImp->evictImage(device, image), ZE_RESULT_SUCCESS);

    Image::fromHandle(image)->destroy();

    contextImp->destroy();
}

HWTEST_F(ContextTest, givenMakeImageResidentThenMakeImageResidentIsCalledWithForcePagingFenceTrue) {
    if (!device->getNEODevice()->getRootDeviceEnvironment().getReleaseHelper() ||
        !device->getNEODevice()->getDeviceInfo().imageSupport) {
        GTEST_SKIP();
    }

    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    class MockMemoryOperationsForcePagingFenceCheck : public NEO::MockMemoryOperations {
      public:
        MockMemoryOperationsForcePagingFenceCheck() {}

        MemoryOperationsStatus makeResident(NEO::Device *neoDevice, ArrayRef<NEO::GraphicsAllocation *> gfxAllocations, bool isDummyExecNeeded, bool forcePagingFence) override {
            makeResidentCalledCount++;
            forcePagingFencePassed = forcePagingFence;
            return MemoryOperationsStatus::success;
        }
        bool forcePagingFencePassed = false;
    };

    auto mockMemoryOperationsInterface = new MockMemoryOperationsForcePagingFenceCheck();
    mockMemoryOperationsInterface->captureGfxAllocationsForMakeResident = true;
    device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[0]->memoryOperationsInterface.reset(mockMemoryOperationsInterface);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    ze_image_handle_t image = {};
    ze_image_desc_t imageDesc = {};
    imageDesc.stype = ZE_STRUCTURE_TYPE_IMAGE_DESC;

    res = contextImp->createImage(device, &imageDesc, &image);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, image);

    contextImp->makeImageResident(device, image);
    EXPECT_EQ(2, mockMemoryOperationsInterface->makeResidentCalledCount);
    EXPECT_TRUE(mockMemoryOperationsInterface->forcePagingFencePassed);

    Image::fromHandle(image)->destroy();

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

class ContextWhiteboxForIpcTesting : public ::L0::ContextImp {
  public:
    ContextWhiteboxForIpcTesting(L0::DriverHandle *driverHandle) : L0::ContextImp(driverHandle) {}

    using ::L0::ContextImp::setIPCHandleData;
};

TEST_F(ContextTest, whenCallingSetIPCHandleDataWithIpcMemoryDataTypeThenIpcDataIsSetInHandleTracking) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextWhiteboxForIpcTesting contextWhitebox(driverHandle.get());

    // Create a mock graphics allocation
    NEO::MockGraphicsAllocation mockAllocation;

    // Set up test data
    uint64_t handle = 12345;
    L0::IpcMemoryData ipcData;
    ipcData.handle = handle;
    ipcData.type = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);
    ipcData.poolOffset = 0;

    uint64_t ptrAddress = 0x1000;
    uint8_t type = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);

    // Verify IPC handle map is initially empty
    EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());

    // Call setIPCHandleData with IpcMemoryData type
    contextWhitebox.setIPCHandleData<L0::IpcMemoryData>(&mockAllocation, handle, ipcData, ptrAddress, type, nullptr, L0::IpcHandleType::fdHandle);

    // Verify the handle was added to the IPC handle map
    auto &ipcHandleMap = driverHandle->getIPCHandleMap();
    EXPECT_EQ(1u, ipcHandleMap.size());

    // Verify the handle tracking entry
    auto handleIterator = ipcHandleMap.find(handle);
    ASSERT_NE(handleIterator, ipcHandleMap.end());

    L0::IpcHandleTracking *handleTracking = handleIterator->second;
    EXPECT_NE(nullptr, handleTracking);
    EXPECT_EQ(&mockAllocation, handleTracking->alloc);
    EXPECT_EQ(1u, handleTracking->refcnt);
    EXPECT_EQ(ptrAddress, handleTracking->ptr);
    EXPECT_EQ(handle, handleTracking->handle);

    // Verify that the ipcData field is set
    uint64_t handleValue = 0;
    memcpy(&handleValue, &handleTracking->ipcData.handle, sizeof(handleValue));
    EXPECT_EQ(handle, handleValue);

    uint8_t typeValue = 0;
    memcpy(&typeValue, &handleTracking->ipcData.type, sizeof(typeValue));
    EXPECT_EQ(type, typeValue);

    // Clean up - remove the entry from the map to avoid issues in teardown
    delete handleTracking;
    driverHandle->getIPCHandleMap().clear();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingSetIPCHandleDataWithIpcOpaqueMemoryDataTypeThenIpcDataIsNotSetInHandleTracking) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextWhiteboxForIpcTesting contextWhitebox(driverHandle.get());

    // Create a mock graphics allocation
    NEO::MockGraphicsAllocation mockAllocation;

    // Set up test data for IpcOpaqueMemoryData
    uint64_t handle = 67890;
    L0::IpcOpaqueMemoryData opaqueIpcData;
    opaqueIpcData.handle.fd = static_cast<int>(handle);
    opaqueIpcData.memoryType = static_cast<uint8_t>(InternalMemoryType::sharedUnifiedMemory);
    opaqueIpcData.processId = 1234;
    opaqueIpcData.type = L0::IpcHandleType::fdHandle;
    opaqueIpcData.poolOffset = 0;

    uint64_t ptrAddress = 0x2000;
    uint8_t type = static_cast<uint8_t>(InternalMemoryType::sharedUnifiedMemory);

    // Verify IPC handle map is initially empty
    EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());

    // Call setIPCHandleData with IpcOpaqueMemoryData type
    contextWhitebox.setIPCHandleData<L0::IpcOpaqueMemoryData>(&mockAllocation, handle, opaqueIpcData, ptrAddress, type, nullptr, L0::IpcHandleType::fdHandle);

    // Verify the handle was added to the IPC handle map
    auto &ipcHandleMap = driverHandle->getIPCHandleMap();
    EXPECT_EQ(1u, ipcHandleMap.size());

    // Verify the handle tracking entry
    auto handleIterator = ipcHandleMap.find(handle);
    ASSERT_NE(handleIterator, ipcHandleMap.end());

    L0::IpcHandleTracking *handleTracking = handleIterator->second;
    EXPECT_NE(nullptr, handleTracking);
    EXPECT_EQ(&mockAllocation, handleTracking->alloc);
    EXPECT_EQ(1u, handleTracking->refcnt);
    EXPECT_EQ(ptrAddress, handleTracking->ptr);
    EXPECT_EQ(handle, handleTracking->handle);

    uint64_t handleTrackingFdValue = 0;
    memcpy(&handleTrackingFdValue, &handleTracking->opaqueData.handle.fd, sizeof(handleTrackingFdValue));
    uint64_t fdValue = 0;
    memcpy(&fdValue, &opaqueIpcData.handle.fd, sizeof(fdValue));
    EXPECT_EQ(fdValue, handleTrackingFdValue);

    uint8_t memoryTypeValue = 0;
    memcpy(&memoryTypeValue, &opaqueIpcData.memoryType, sizeof(memoryTypeValue));
    uint8_t memoryTypeHandleValue = 0;
    memcpy(&memoryTypeHandleValue, &handleTracking->opaqueData.memoryType, sizeof(memoryTypeHandleValue));
    EXPECT_EQ(memoryTypeValue, memoryTypeHandleValue);

    unsigned int processIdHandleValue = 0;
    memcpy(&processIdHandleValue, &handleTracking->opaqueData.processId, sizeof(processIdHandleValue));
    unsigned int processIdValue = 0;
    memcpy(&processIdValue, &opaqueIpcData.processId, sizeof(processIdValue));
    EXPECT_EQ(processIdValue, processIdHandleValue);

    EXPECT_EQ(opaqueIpcData.type, handleTracking->opaqueData.type);

    uint64_t poolOffsetHandleValue = 0;
    memcpy(&poolOffsetHandleValue, &handleTracking->opaqueData.poolOffset, sizeof(poolOffsetHandleValue));
    uint64_t poolOffsetValue = 0;
    memcpy(&poolOffsetValue, &opaqueIpcData.poolOffset, sizeof(poolOffsetValue));
    EXPECT_EQ(poolOffsetValue, poolOffsetHandleValue);

    // Clean up - remove the entry from the map to avoid issues in teardown
    delete handleTracking;
    driverHandle->getIPCHandleMap().clear();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingSetIPCHandleDataWithInvalidIpcHandleTypeThenHandleUnionIsNotSet) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextWhiteboxForIpcTesting contextWhitebox(driverHandle.get());

    // Create a mock graphics allocation
    NEO::MockGraphicsAllocation mockAllocation;

    // Set up test data with invalid IpcHandleType
    uint64_t handle = 11111;
    L0::IpcOpaqueMemoryData opaqueIpcData;
    uint64_t ptrAddress = 0x3000;
    uint8_t type = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);

    // Test with maxHandle (invalid but defined)
    L0::IpcHandleType invalidHandleType1 = L0::IpcHandleType::maxHandle;

    // Verify IPC handle map is initially empty
    EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());

    // Call setIPCHandleData with invalid handle type (maxHandle)
    contextWhitebox.setIPCHandleData<L0::IpcOpaqueMemoryData>(&mockAllocation, handle, opaqueIpcData, ptrAddress, type, nullptr, invalidHandleType1);

    // Verify the handle was added to the IPC handle map
    auto &ipcHandleMap = driverHandle->getIPCHandleMap();
    EXPECT_EQ(1u, ipcHandleMap.size());

    // Verify the IpcOpaqueMemoryData values
    EXPECT_EQ(type, opaqueIpcData.memoryType);
    EXPECT_EQ(invalidHandleType1, opaqueIpcData.type);
    EXPECT_NE(0u, opaqueIpcData.processId); // Should be set to current process ID

    auto handleIterator = ipcHandleMap.find(handle);
    ASSERT_NE(handleIterator, ipcHandleMap.end());
    delete handleIterator->second;
    driverHandle->getIPCHandleMap().clear();

    // Test with completely invalid handle type (beyond enum range)
    handle = 22222;
    opaqueIpcData = {}; // Reset
    L0::IpcHandleType invalidHandleType2 = L0::IpcHandleType::maxHandle;

    contextWhitebox.setIPCHandleData<L0::IpcOpaqueMemoryData>(&mockAllocation, handle, opaqueIpcData, ptrAddress, type, nullptr, invalidHandleType2);

    // Verify the handle was added to the IPC handle map
    EXPECT_EQ(1u, ipcHandleMap.size());

    // Verify the IpcOpaqueMemoryData values
    EXPECT_EQ(type, opaqueIpcData.memoryType);
    EXPECT_EQ(invalidHandleType2, opaqueIpcData.type);
    EXPECT_NE(0u, opaqueIpcData.processId); // Should be set to current process ID

    handleIterator = ipcHandleMap.find(handle);
    ASSERT_NE(handleIterator, ipcHandleMap.end());
    delete handleIterator->second;
    driverHandle->getIPCHandleMap().clear();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingSetIPCHandleDataWithNtHandleTypeThenHandleUnionIsSetCorrectly) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextWhiteboxForIpcTesting contextWhitebox(driverHandle.get());

    // Create a mock graphics allocation
    NEO::MockGraphicsAllocation mockAllocation;

    // Set up test data for ntHandle
    uint64_t handle = 98765;
    L0::IpcOpaqueMemoryData opaqueIpcData;
    opaqueIpcData.handle.reserved = handle;
    opaqueIpcData.memoryType = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);
    opaqueIpcData.processId = 5678;
    opaqueIpcData.type = L0::IpcHandleType::ntHandle;
    opaqueIpcData.poolOffset = 0;

    uint64_t ptrAddress = 0x4000;
    uint8_t type = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);

    // Verify IPC handle map is initially empty
    EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());

    // Call setIPCHandleData with ntHandle type
    contextWhitebox.setIPCHandleData<L0::IpcOpaqueMemoryData>(&mockAllocation, handle, opaqueIpcData, ptrAddress, type, nullptr, L0::IpcHandleType::ntHandle);

    // Verify the handle was added to the IPC handle map
    auto &ipcHandleMap = driverHandle->getIPCHandleMap();
    EXPECT_EQ(1u, ipcHandleMap.size());

    // Verify the handle tracking entry
    auto handleIterator = ipcHandleMap.find(handle);
    ASSERT_NE(handleIterator, ipcHandleMap.end());

    L0::IpcHandleTracking *handleTracking = handleIterator->second;
    EXPECT_NE(nullptr, handleTracking);
    EXPECT_EQ(&mockAllocation, handleTracking->alloc);
    EXPECT_EQ(1u, handleTracking->refcnt);
    EXPECT_EQ(ptrAddress, handleTracking->ptr);
    EXPECT_EQ(handle, handleTracking->handle);

    uint64_t handleTrackingReservedValue = 0;
    memcpy(&handleTrackingReservedValue, &handleTracking->opaqueData.handle.reserved, sizeof(handleTrackingReservedValue));
    uint64_t reservedDataValue = 0;
    memcpy(&reservedDataValue, &opaqueIpcData.handle.reserved, sizeof(reservedDataValue));
    EXPECT_EQ(reservedDataValue, handleTrackingReservedValue);

    uint8_t memoryTypeValue = 0;
    memcpy(&memoryTypeValue, &opaqueIpcData.memoryType, sizeof(memoryTypeValue));
    uint8_t memoryTypeHandleValue = 0;
    memcpy(&memoryTypeHandleValue, &handleTracking->opaqueData.memoryType, sizeof(memoryTypeHandleValue));
    EXPECT_EQ(memoryTypeValue, memoryTypeHandleValue);

    unsigned int processIdHandleValue = 0;
    memcpy(&processIdHandleValue, &handleTracking->opaqueData.processId, sizeof(processIdHandleValue));
    unsigned int processIdValue = 0;
    memcpy(&processIdValue, &opaqueIpcData.processId, sizeof(processIdValue));
    EXPECT_EQ(processIdValue, processIdHandleValue);

    EXPECT_EQ(opaqueIpcData.type, handleTracking->opaqueData.type);

    uint64_t poolOffsetHandleValue = 0;
    memcpy(&poolOffsetHandleValue, &handleTracking->opaqueData.poolOffset, sizeof(poolOffsetHandleValue));
    uint64_t poolOffsetValue = 0;
    memcpy(&poolOffsetValue, &opaqueIpcData.poolOffset, sizeof(poolOffsetValue));
    EXPECT_EQ(poolOffsetValue, poolOffsetHandleValue);

    // Verify the IpcOpaqueMemoryData values are set correctly
    EXPECT_EQ(type, opaqueIpcData.memoryType);
    EXPECT_EQ(L0::IpcHandleType::ntHandle, opaqueIpcData.type);
    uint64_t reservedValue = 0;
    memcpy(&reservedValue, &opaqueIpcData.handle.reserved, sizeof(reservedValue));
    EXPECT_EQ(handle, reservedValue);
    unsigned int processID = 0;
    memcpy(&processID, &opaqueIpcData.processId, sizeof(processID));
    EXPECT_NE(0u, processID); // Should be set to current process ID

    // Clean up - remove the entry from the map to avoid issues in teardown
    delete handleTracking;
    driverHandle->getIPCHandleMap().clear();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingGetFdFromIpcHandleWithOpaqueFdHandleThenCorrectHandleIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ::L0::ContextImp contextWhitebox{driverHandle.get()};

    // Configure for opaque handle with fd type
    contextWhitebox.settings.useOpaqueHandle = true;
    contextWhitebox.settings.handleType = IpcHandleType::fdHandle;

    NEO::MockGraphicsAllocation mockAllocation;
    uint64_t testHandle = 12345;
    int testFd = 42;

    // Create IPC handle tracking entry for opaque fd handle
    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = &mockAllocation;
    handleTracking->refcnt = 1;
    handleTracking->ptr = 0x1000;
    handleTracking->handle = testHandle;
    handleTracking->opaqueData.handle.fd = testFd;
    handleTracking->opaqueData.type = IpcHandleType::fdHandle;

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testFd] = handleTracking;
    }

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcOpaqueMemoryData *opaqueData = reinterpret_cast<L0::IpcOpaqueMemoryData *>(ipcHandle.data);
    opaqueData->type = IpcHandleType::fdHandle;
    opaqueData->handle.fd = testFd;

    // Call getFdFromIpcHandle
    uint64_t returnedHandle = 0;
    res = contextWhitebox.getFdFromIpcHandle(ipcHandle, &returnedHandle);

    // Verify success and correct handle returned (should return the key, which is testFd)
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(static_cast<uint64_t>(testFd), returnedHandle);

    delete handleTracking;
    driverHandle->getIPCHandleMap().clear();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingGetFdFromIpcHandleWithNonOpaqueHandleThenCorrectFdIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ::L0::ContextImp contextWhitebox(driverHandle.get());

    // Configure for non-opaque handle
    contextWhitebox.settings.useOpaqueHandle = false;
    contextWhitebox.settings.handleType = IpcHandleType::fdHandle;

    NEO::MockGraphicsAllocation mockAllocation;
    uint64_t testHandle = 67890;

    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = &mockAllocation;
    handleTracking->refcnt = 1;
    handleTracking->ptr = 0x2000;
    handleTracking->handle = testHandle;
    handleTracking->ipcData.handle = testHandle;
    handleTracking->ipcData.type = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testHandle] = handleTracking;
    }

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcMemoryData *memoryData = reinterpret_cast<L0::IpcMemoryData *>(ipcHandle.data);
    memoryData->handle = testHandle;
    memoryData->type = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);

    // Call getFdFromIpcHandle
    uint64_t returnedHandle = 0;
    res = contextWhitebox.getFdFromIpcHandle(ipcHandle, &returnedHandle);

    // Verify success and correct handle returned
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(testHandle, returnedHandle);

    delete handleTracking;
    driverHandle->getIPCHandleMap().clear();

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingGetFdFromIpcHandleWithOpaqueNtHandleThenErrorIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ::L0::ContextImp contextWhitebox(driverHandle.get());

    // Configure for opaque handle with NT type (should fail since getFdFromIpcHandle expects fd)
    contextWhitebox.settings.useOpaqueHandle = true;
    contextWhitebox.settings.handleType = IpcHandleType::ntHandle;

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcOpaqueMemoryData *opaqueData = reinterpret_cast<L0::IpcOpaqueMemoryData *>(ipcHandle.data);
    opaqueData->handle.reserved = 12345;
    opaqueData->type = IpcHandleType::ntHandle;

    // Call getFdFromIpcHandle
    uint64_t returnedHandle = 0;
    res = contextWhitebox.getFdFromIpcHandle(ipcHandle, &returnedHandle);

    // Should return error since NT handle is not supported for fd extraction
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingGetFdFromIpcHandleWithInvalidHandleNotInMapThenErrorIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ::L0::ContextImp contextWhitebox(driverHandle.get());

    // Configure for opaque handle with fd type
    contextWhitebox.settings.useOpaqueHandle = true;
    contextWhitebox.settings.handleType = IpcHandleType::fdHandle;

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcOpaqueMemoryData *opaqueData = reinterpret_cast<L0::IpcOpaqueMemoryData *>(ipcHandle.data);
    opaqueData->handle.fd = 999; // Handle not in map
    opaqueData->type = IpcHandleType::fdHandle;

    // Call getFdFromIpcHandle
    uint64_t returnedHandle = 0;
    res = contextWhitebox.getFdFromIpcHandle(ipcHandle, &returnedHandle);

    // Should return error since handle is not found in map
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingGetFdFromIpcHandleWithNonOpaqueHandleNotInMapThenErrorIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ::L0::ContextImp contextWhitebox(driverHandle.get());

    // Configure for non-opaque handle
    contextWhitebox.settings.useOpaqueHandle = false;
    contextWhitebox.settings.handleType = IpcHandleType::fdHandle;

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcMemoryData *memoryData = reinterpret_cast<L0::IpcMemoryData *>(ipcHandle.data);
    memoryData->handle = 888; // Handle not in map
    memoryData->type = static_cast<uint8_t>(InternalMemoryType::deviceUnifiedMemory);

    // Call getFdFromIpcHandle
    uint64_t returnedHandle = 0;
    res = contextWhitebox.getFdFromIpcHandle(ipcHandle, &returnedHandle);

    // Should return error since handle is not found in map
    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingPutIpcMemHandleWithOpaqueIpcHandleThenUnregisterIpcHandleWithServerIsCalled) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    contextImp->settings.useOpaqueHandle = true;
    contextImp->settings.handleType = IpcHandleType::fdHandle;

    NEO::MockGraphicsAllocation mockAllocation;
    uint64_t testHandle = 12345;

    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = &mockAllocation;
    handleTracking->refcnt = 1;
    handleTracking->ptr = 0x1000;
    handleTracking->handle = testHandle;
    handleTracking->opaqueData.handle.fd = static_cast<int>(testHandle);
    handleTracking->opaqueData.type = IpcHandleType::fdHandle;

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testHandle] = handleTracking;
    }

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcMemoryData *ipcData = reinterpret_cast<L0::IpcMemoryData *>(ipcHandle.data);
    ipcData->handle = testHandle;

    res = contextImp->putIpcMemHandle(ipcHandle);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingPutIpcMemHandleWithNonOpaqueIpcHandleThenUnregisterIpcHandleWithServerIsNotCalled) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    contextImp->settings.useOpaqueHandle = false;

    NEO::MockGraphicsAllocation mockAllocation;
    uint64_t testHandle = 67890;

    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = &mockAllocation;
    handleTracking->refcnt = 1;
    handleTracking->ptr = 0x2000;
    handleTracking->handle = testHandle;
    handleTracking->ipcData.handle = testHandle;

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testHandle] = handleTracking;
    }

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcMemoryData *ipcData = reinterpret_cast<L0::IpcMemoryData *>(ipcHandle.data);
    ipcData->handle = testHandle;

    res = contextImp->putIpcMemHandle(ipcHandle);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingPutIpcMemHandleWithMultipleRefCountThenHandleNotRemovedUntilLastRef) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    contextImp->settings.useOpaqueHandle = true;
    contextImp->settings.handleType = IpcHandleType::fdHandle;

    NEO::MockGraphicsAllocation mockAllocation;
    uint64_t testHandle = 11111;

    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = &mockAllocation;
    handleTracking->refcnt = 2;
    handleTracking->ptr = 0x3000;
    handleTracking->handle = testHandle;
    handleTracking->opaqueData.handle.fd = static_cast<int>(testHandle);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testHandle] = handleTracking;
    }

    ze_ipc_mem_handle_t ipcHandle = {};
    L0::IpcMemoryData *ipcData = reinterpret_cast<L0::IpcMemoryData *>(ipcHandle.data);
    ipcData->handle = testHandle;

    res = contextImp->putIpcMemHandle(ipcHandle);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        EXPECT_FALSE(driverHandle->getIPCHandleMap().empty());
        auto it = driverHandle->getIPCHandleMap().find(testHandle);
        ASSERT_NE(it, driverHandle->getIPCHandleMap().end());
        EXPECT_EQ(1u, it->second->refcnt);
    }

    res = contextImp->putIpcMemHandle(ipcHandle);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

class ZexMemFreeRegisterCallbackExtTests : public Test<DeviceFixture> {
  public:
    void SetUp() override {
        Test<DeviceFixture>::SetUp();
        testCallbackExecuted = false;
        testCallbackUserData = nullptr;
    }

    void TearDown() override {
        testCallbackExecuted = false;
        testCallbackUserData = nullptr;
        Test<DeviceFixture>::TearDown();
    }

    // Test callback function
    static void testMemoryFreeCallback(void *pUserData) {
        testCallbackExecuted = true;
        testCallbackUserData = pUserData;
    }

    static bool testCallbackExecuted;
    static void *testCallbackUserData;
};

// Static member definitions
bool ZexMemFreeRegisterCallbackExtTests::testCallbackExecuted = false;
void *ZexMemFreeRegisterCallbackExtTests::testCallbackUserData = nullptr;

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithValidParametersThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Allocate memory
    const size_t allocSize = 4096;
    void *ptr = nullptr;
    ze_host_mem_alloc_desc_t hostDesc = {ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC};
    res = zeMemAllocHost(hContext, &hostDesc, allocSize, 0u, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    // Setup callback descriptor
    const char *userData = "zex api test data";
    zex_memory_free_callback_ext_desc_t callbackDesc = {};
    callbackDesc.stype = ZEX_STRUCTURE_TYPE_MEMORY_FREE_CALLBACK_EXT_DESC;
    callbackDesc.pNext = nullptr;
    callbackDesc.pfnCallback = testMemoryFreeCallback;
    callbackDesc.pUserData = const_cast<char *>(userData);

    // Call the ZEX API function
    res = zexMemFreeRegisterCallbackExt(hContext, &callbackDesc, ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = zeMemFree(hContext, ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(true, testCallbackExecuted);
    EXPECT_EQ(testCallbackUserData, const_cast<char *>(userData));

    res = zeContextDestroy(hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithNullContextThenInvalidArgumentIsReturned) {
    // Setup callback descriptor
    const char *userData = "test data";
    zex_memory_free_callback_ext_desc_t callbackDesc = {};
    callbackDesc.stype = ZEX_STRUCTURE_TYPE_MEMORY_FREE_CALLBACK_EXT_DESC;
    callbackDesc.pNext = nullptr;
    callbackDesc.pfnCallback = testMemoryFreeCallback;
    callbackDesc.pUserData = const_cast<char *>(userData);

    void *ptr = reinterpret_cast<void *>(0x1234);

    // Call with null context
    ze_result_t res = zexMemFreeRegisterCallbackExt(nullptr, &callbackDesc, ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);
}

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithNullDescriptorThenInvalidArgumentIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    void *ptr = reinterpret_cast<void *>(0x1234);

    // Call with null descriptor
    res = zexMemFreeRegisterCallbackExt(hContext, nullptr, ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = zeContextDestroy(hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithNullPointerThenInvalidArgumentIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Setup callback descriptor
    const char *userData = "test data";
    zex_memory_free_callback_ext_desc_t callbackDesc = {};
    callbackDesc.stype = ZEX_STRUCTURE_TYPE_MEMORY_FREE_CALLBACK_EXT_DESC;
    callbackDesc.pNext = nullptr;
    callbackDesc.pfnCallback = testMemoryFreeCallback;
    callbackDesc.pUserData = const_cast<char *>(userData);

    // Call with null pointer
    res = zexMemFreeRegisterCallbackExt(hContext, &callbackDesc, nullptr);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = zeContextDestroy(hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithInvalidPointerThenInvalidArgumentIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Setup callback descriptor
    const char *userData = "test data";
    zex_memory_free_callback_ext_desc_t callbackDesc = {};
    callbackDesc.stype = ZEX_STRUCTURE_TYPE_MEMORY_FREE_CALLBACK_EXT_DESC;
    callbackDesc.pNext = nullptr;
    callbackDesc.pfnCallback = testMemoryFreeCallback;
    callbackDesc.pUserData = const_cast<char *>(userData);

    // Call with invalid pointer (not an SVM allocation)
    void *invalidPtr = reinterpret_cast<void *>(0x1234);
    res = zexMemFreeRegisterCallbackExt(hContext, &callbackDesc, invalidPtr);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = zeContextDestroy(hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithDeviceMemoryThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    const size_t allocSize = 4096;
    void *ptr = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC};
    res = zeMemAllocDevice(hContext, &deviceDesc, allocSize, 0u, device->toHandle(), &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    // Setup callback descriptor
    const char *userData = "device memory test";
    zex_memory_free_callback_ext_desc_t callbackDesc = {};
    callbackDesc.stype = ZEX_STRUCTURE_TYPE_MEMORY_FREE_CALLBACK_EXT_DESC;
    callbackDesc.pNext = nullptr;
    callbackDesc.pfnCallback = testMemoryFreeCallback;
    callbackDesc.pUserData = const_cast<char *>(userData);

    // Call the ZEX API function
    res = zexMemFreeRegisterCallbackExt(hContext, &callbackDesc, ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = zeMemFree(hContext, ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(true, testCallbackExecuted);
    EXPECT_EQ(testCallbackUserData, const_cast<char *>(userData));

    res = zeContextDestroy(hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithSharedMemoryThenSuccessIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Allocate shared memory
    const size_t allocSize = 4096;
    void *ptr = nullptr;
    ze_device_mem_alloc_desc_t deviceDesc = {ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC};
    ze_host_mem_alloc_desc_t hostDesc = {ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC};
    res = zeMemAllocShared(hContext, &deviceDesc, &hostDesc, allocSize, 0u, device->toHandle(), &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    // Setup callback descriptor
    const char *userData = "shared memory test";
    zex_memory_free_callback_ext_desc_t callbackDesc = {};
    callbackDesc.stype = ZEX_STRUCTURE_TYPE_MEMORY_FREE_CALLBACK_EXT_DESC;
    callbackDesc.pNext = nullptr;
    callbackDesc.pfnCallback = testMemoryFreeCallback;
    callbackDesc.pUserData = const_cast<char *>(userData);

    // Call the ZEX API function
    res = zexMemFreeRegisterCallbackExt(hContext, &callbackDesc, ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = zeMemFree(hContext, ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(true, testCallbackExecuted);
    EXPECT_EQ(testCallbackUserData, const_cast<char *>(userData));

    res = zeContextDestroy(hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ZexMemFreeRegisterCallbackExtTests, whenCallingZexMemFreeRegisterCallbackExtWithNullCallbackFunctionThenInvalidArgumentIsReturned) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Allocate memory
    const size_t allocSize = 4096;
    void *ptr = nullptr;
    ze_host_mem_alloc_desc_t hostDesc = {ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC};
    res = zeMemAllocHost(hContext, &hostDesc, allocSize, 0u, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    // Setup callback descriptor with null callback function
    const char *userData = "null callback test";
    zex_memory_free_callback_ext_desc_t callbackDesc = {};
    callbackDesc.stype = ZEX_STRUCTURE_TYPE_MEMORY_FREE_CALLBACK_EXT_DESC;
    callbackDesc.pNext = nullptr;
    callbackDesc.pfnCallback = nullptr; // null callback function
    callbackDesc.pUserData = const_cast<char *>(userData);

    // Call the ZEX API function - this should fail due to the invalid callback function pointer
    res = zexMemFreeRegisterCallbackExt(hContext, &callbackDesc, ptr);
    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);

    res = zeMemFree(hContext, ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    res = zeContextDestroy(hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenSettingVirtualMemAccessAttributeWithChangedFlagsAndMappedAllocationsThenUnmapAndRemapIsCalled) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    // Reserve virtual memory
    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Create physical memory and map it
    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = contextImp->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;
    res = contextImp->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Change access attribute to trigger unmap/remap logic
    ze_memory_access_attribute_t newAccess = {ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY};
    res = contextImp->setVirtualMemAccessAttribute(ptr, pagesize, newAccess);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    // Check that the new access attribute is set
    ze_memory_access_attribute_t outAccess = {};
    size_t outSize = 0;
    res = contextImp->getVirtualMemAccessAttribute(ptr, pagesize, &outAccess, &outSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(outAccess, newAccess);

    res = contextImp->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, givenContextWhenChangingVirtualMemAccessAttributeThenUnmapAndMapAreCalled) {
    auto ctx = std::make_unique<Mock<ContextImp>>(driverHandle.get());
    ctx->createPhysicalMemCallBase = true;
    ctx->queryVirtualMemPageSizeCallBase = true;
    ctx->reserveVirtualMemCallBase = true;
    ctx->queryVirtualMemPageSizeWithStartAddressCallBase = true;
    ctx->mapVirtualMemCallBase = true;
    ctx->setVirtualMemAccessAttributeCallBase = true;
    ctx->unMapVirtualMemCallBase = true;
    ctx->destroyPhysicalMemCallBase = true;
    ctx->freeVirtualMemCallBase = true;

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    ze_result_t res = ctx->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = ctx->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    size_t offset = 0;
    res = ctx->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t newAccess = {ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY};
    res = ctx->setVirtualMemAccessAttribute(ptr, pagesize, newAccess);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(ctx->mapVirtualMemCalled, 2u);
    EXPECT_EQ(ctx->unMapVirtualMemCalled, 1u);

    res = ctx->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, givenContextWhenChangingVirtualMemAccessAttributeToNoneThenUnmapAndMapAreCalled) {
    auto ctx = std::make_unique<Mock<ContextImp>>(driverHandle.get());
    ctx->createPhysicalMemCallBase = true;
    ctx->queryVirtualMemPageSizeCallBase = true;
    ctx->reserveVirtualMemCallBase = true;
    ctx->queryVirtualMemPageSizeWithStartAddressCallBase = true;
    ctx->mapVirtualMemCallBase = true;
    ctx->setVirtualMemAccessAttributeCallBase = true;
    ctx->unMapVirtualMemCallBase = true;
    ctx->destroyPhysicalMemCallBase = true;
    ctx->freeVirtualMemCallBase = true;

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    ze_result_t res = ctx->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = ctx->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY};
    size_t offset = 0;
    res = ctx->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t newAccess = {ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};
    res = ctx->setVirtualMemAccessAttribute(ptr, pagesize, newAccess);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(ctx->mapVirtualMemCalled, 2u);
    EXPECT_EQ(ctx->unMapVirtualMemCalled, 1u);

    res = ctx->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, givenContextWhenChangingVirtualMemAccessAttributeToReadWriteThenUnmapAndMapAreCalled) {
    auto ctx = std::make_unique<Mock<ContextImp>>(driverHandle.get());
    ctx->createPhysicalMemCallBase = true;
    ctx->queryVirtualMemPageSizeCallBase = true;
    ctx->reserveVirtualMemCallBase = true;
    ctx->queryVirtualMemPageSizeWithStartAddressCallBase = true;
    ctx->mapVirtualMemCallBase = true;
    ctx->setVirtualMemAccessAttributeCallBase = true;
    ctx->unMapVirtualMemCallBase = true;
    ctx->destroyPhysicalMemCallBase = true;
    ctx->freeVirtualMemCallBase = true;

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pagesize = 0u;
    ze_result_t res = ctx->queryVirtualMemPageSize(device, size, &pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->reserveVirtualMem(pStart, pagesize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pagesize};
    ze_physical_mem_handle_t mem = {};
    res = ctx->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_NONE};
    size_t offset = 0;
    res = ctx->mapVirtualMem(ptr, pagesize, mem, offset, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    NEO::VirtualMemoryReservation *virtualMemoryReservation = nullptr;
    auto lockVirtual = this->driverHandle->getMemoryManager()->lockVirtualMemoryReservationMap();
    virtualMemoryReservation = ctx->findSupportedVirtualReservation(ptr, size);
    virtualMemoryReservation->flags.readOnly = false;
    virtualMemoryReservation->flags.noAccess = false;
    virtualMemoryReservation->flags.readWrite = false;
    lockVirtual.unlock();
    ze_memory_access_attribute_t newAccess = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    res = ctx->setVirtualMemAccessAttribute(ptr, pagesize, newAccess);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(ctx->mapVirtualMemCalled, 2u);
    EXPECT_EQ(ctx->unMapVirtualMemCalled, 1u);

    res = ctx->unMapVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->freeVirtualMem(ptr, pagesize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, givenContextWhenChangingVirtualMemAccessAttributeAndNoMappedAllocationsThenUnmapAndMapAreNotCalled) {
    auto ctx = std::make_unique<Mock<ContextImp>>(driverHandle.get());
    ctx->queryVirtualMemPageSizeCallBase = true;
    ctx->reserveVirtualMemCallBase = true;
    ctx->setVirtualMemAccessAttributeCallBase = true;
    ctx->freeVirtualMemCallBase = true;
    ctx->queryVirtualMemPageSizeWithStartAddressCallBase = true;

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pageSize = 0u;

    ze_result_t res = ctx->queryVirtualMemPageSize(device, size, &pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->reserveVirtualMem(pStart, pageSize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t newAccess = {ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY};
    res = ctx->setVirtualMemAccessAttribute(ptr, pageSize, newAccess);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(ctx->mapVirtualMemCalled, 0u);
    EXPECT_EQ(ctx->unMapVirtualMemCalled, 0u);

    res = ctx->freeVirtualMem(ptr, pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, givenContextWhenChangingVirtualMemAccessAttributeAndFlagsUnchangedThenUnmapAndMapAreNotCalled) {
    auto ctx = std::make_unique<Mock<ContextImp>>(driverHandle.get());
    ctx->queryVirtualMemPageSizeCallBase = true;
    ctx->reserveVirtualMemCallBase = true;
    ctx->createPhysicalMemCallBase = true;
    ctx->mapVirtualMemCallBase = true;
    ctx->setVirtualMemAccessAttributeCallBase = true;
    ctx->unMapVirtualMemCallBase = true;
    ctx->destroyPhysicalMemCallBase = true;
    ctx->freeVirtualMemCallBase = true;
    ctx->queryVirtualMemPageSizeWithStartAddressCallBase = true;

    void *pStart = 0x0;
    size_t size = 4096u;
    void *ptr = nullptr;
    size_t pageSize = 0u;

    ze_result_t res = ctx->queryVirtualMemPageSize(device, size, &pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->reserveVirtualMem(pStart, pageSize, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pageSize};
    ze_physical_mem_handle_t mem = {};
    res = ctx->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    res = ctx->mapVirtualMem(ptr, pageSize, mem, 0, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(ctx->mapVirtualMemCalled, 1u);

    res = ctx->setVirtualMemAccessAttribute(ptr, pageSize, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(ctx->mapVirtualMemCalled, 1u);
    EXPECT_EQ(ctx->unMapVirtualMemCalled, 0u);

    res = ctx->unMapVirtualMem(ptr, pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->freeVirtualMem(ptr, pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, givenContextWhenChangingVirtualMemAccessAttributeAndPtrNotInMappedAllocationsThenUnmapAndMapAreNotCalled) {
    auto ctx = std::make_unique<Mock<ContextImp>>(driverHandle.get());
    ctx->queryVirtualMemPageSizeCallBase = true;
    ctx->reserveVirtualMemCallBase = true;
    ctx->createPhysicalMemCallBase = true;
    ctx->mapVirtualMemCallBase = true;
    ctx->setVirtualMemAccessAttributeCallBase = true;
    ctx->unMapVirtualMemCallBase = true;
    ctx->destroyPhysicalMemCallBase = true;
    ctx->freeVirtualMemCallBase = true;
    ctx->queryVirtualMemPageSizeWithStartAddressCallBase = true;

    void *pStart = 0x0;
    size_t totalSize = 8192u;
    void *reservationBase = nullptr;
    size_t pageSize = 0u;

    ze_result_t res = ctx->queryVirtualMemPageSize(device, totalSize, &pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->reserveVirtualMem(pStart, pageSize * 2, &reservationBase);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    void *secondPagePtr = reinterpret_cast<void *>(reinterpret_cast<uint64_t>(reservationBase) + pageSize);
    ze_physical_mem_desc_t descMem = {ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC, nullptr, 0, pageSize};
    ze_physical_mem_handle_t mem = {};
    res = ctx->createPhysicalMem(device, &descMem, &mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ze_memory_access_attribute_t access = {ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE};
    res = ctx->mapVirtualMem(secondPagePtr, pageSize, mem, 0, access);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_EQ(ctx->mapVirtualMemCalled, 1u);

    ze_memory_access_attribute_t newAccess = {ZE_MEMORY_ACCESS_ATTRIBUTE_READONLY};
    res = ctx->setVirtualMemAccessAttribute(reservationBase, pageSize, newAccess);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    EXPECT_EQ(ctx->mapVirtualMemCalled, 1u);
    EXPECT_EQ(ctx->unMapVirtualMemCalled, 0u);

    res = ctx->unMapVirtualMem(secondPagePtr, pageSize);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->destroyPhysicalMem(mem);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    res = ctx->freeVirtualMem(reservationBase, pageSize * 2);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingFreeMemWithOpaqueIpcHandleThenUnregisterIpcHandleWithServerIsCalled) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    contextImp->settings.useOpaqueHandle = true;
    contextImp->settings.handleType = IpcHandleType::fdHandle;

    size_t size = 4096;
    size_t alignment = 4096;
    void *ptr = nullptr;

    ze_device_mem_alloc_desc_t deviceDesc = {};
    deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC;
    res = contextImp->allocDeviceMem(device->toHandle(), &deviceDesc, size, alignment, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    uint64_t testHandle = 12345;
    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = nullptr;
    handleTracking->refcnt = 1;
    handleTracking->ptr = reinterpret_cast<uint64_t>(ptr);
    handleTracking->handle = testHandle;
    handleTracking->opaqueData.handle.fd = static_cast<int>(testHandle);
    handleTracking->opaqueData.type = IpcHandleType::fdHandle;

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testHandle] = handleTracking;
    }

    res = contextImp->freeMem(ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingFreeMemWithNonOpaqueIpcHandleThenUnregisterIpcHandleWithServerIsNotCalled) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    contextImp->settings.useOpaqueHandle = false;
    contextImp->settings.handleType = IpcHandleType::maxHandle;

    size_t size = 4096;
    size_t alignment = 4096;
    void *ptr = nullptr;

    ze_device_mem_alloc_desc_t deviceDesc = {};
    deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC;
    res = contextImp->allocDeviceMem(device->toHandle(), &deviceDesc, size, alignment, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    uint64_t testHandle = 67890;
    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = nullptr;
    handleTracking->refcnt = 1;
    handleTracking->ptr = reinterpret_cast<uint64_t>(ptr);
    handleTracking->handle = testHandle;
    handleTracking->ipcData.handle = testHandle;

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testHandle] = handleTracking;
    }

    res = contextImp->freeMem(ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

TEST_F(ContextTest, whenCallingFreeMemWithOpaqueNtHandleThenUnregisterIpcHandleWithServerIsNotCalled) {
    ze_context_handle_t hContext;
    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};

    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));

    contextImp->settings.useOpaqueHandle = true;
    contextImp->settings.handleType = IpcHandleType::ntHandle;

    size_t size = 4096;
    size_t alignment = 4096;
    void *ptr = nullptr;

    ze_device_mem_alloc_desc_t deviceDesc = {};
    deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC;
    res = contextImp->allocDeviceMem(device->toHandle(), &deviceDesc, size, alignment, &ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
    EXPECT_NE(nullptr, ptr);

    uint64_t testHandle = 11111;
    L0::IpcHandleTracking *handleTracking = new L0::IpcHandleTracking();
    handleTracking->alloc = nullptr;
    handleTracking->refcnt = 1;
    handleTracking->ptr = reinterpret_cast<uint64_t>(ptr);
    handleTracking->handle = testHandle;
    handleTracking->opaqueData.handle.reserved = testHandle;
    handleTracking->opaqueData.type = IpcHandleType::ntHandle;

    {
        auto lock = driverHandle->lockIPCHandleMap();
        driverHandle->getIPCHandleMap()[testHandle] = handleTracking;
    }

    res = contextImp->freeMem(ptr);
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);

    {
        auto lock = driverHandle->lockIPCHandleMap();
        EXPECT_TRUE(driverHandle->getIPCHandleMap().empty());
    }

    res = contextImp->destroy();
    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
}

} // namespace ult
} // namespace L0
