diff --git a/level_zero/core/test/unit_tests/mock.h b/level_zero/core/test/unit_tests/mock.h
index 938eed0be5..cd2c6f0c1a 100644
--- a/level_zero/core/test/unit_tests/mock.h
+++ b/level_zero/core/test/unit_tests/mock.h
@@ -6,8 +6,6 @@
  */
 
 #pragma once
-#include "gmock/gmock.h"
-
 namespace L0 {
 namespace ult {
 
diff --git a/level_zero/core/test/unit_tests/mocks/CMakeLists.txt b/level_zero/core/test/unit_tests/mocks/CMakeLists.txt
index e4d83cb4fa..b876d15781 100644
--- a/level_zero/core/test/unit_tests/mocks/CMakeLists.txt
+++ b/level_zero/core/test/unit_tests/mocks/CMakeLists.txt
@@ -15,6 +15,7 @@ set(L0_MOCKS_SOURCES
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_cmdlist.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_cmdqueue.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_cmdqueue.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/mock_context.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_device.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_device.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_driver.h
@@ -27,7 +28,6 @@ set(L0_MOCKS_SOURCES
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_kernel.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_l0_debugger.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_memory_manager.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/mock_memory_manager.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_module.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_module.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_sampler.h
diff --git a/level_zero/core/test/unit_tests/mocks/mock_cmdqueue.h b/level_zero/core/test/unit_tests/mocks/mock_cmdqueue.h
index 14df52ccba..9992191edb 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_cmdqueue.h
+++ b/level_zero/core/test/unit_tests/mocks/mock_cmdqueue.h
@@ -11,6 +11,8 @@
 #include "level_zero/core/test/unit_tests/mock.h"
 #include "level_zero/core/test/unit_tests/white_box.h"
 
+#include "gmock/gmock.h"
+
 namespace L0 {
 namespace ult {
 
diff --git a/level_zero/core/test/unit_tests/mocks/mock_context.h b/level_zero/core/test/unit_tests/mocks/mock_context.h
index 0560d9a6df..d56e3e0b29 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_context.h
+++ b/level_zero/core/test/unit_tests/mocks/mock_context.h
@@ -12,6 +12,8 @@
 #include "level_zero/core/test/unit_tests/mock.h"
 #include "level_zero/core/test/unit_tests/white_box.h"
 
+#include "gmock/gmock.h"
+
 namespace L0 {
 namespace ult {
 
diff --git a/level_zero/core/test/unit_tests/mocks/mock_device.h b/level_zero/core/test/unit_tests/mocks/mock_device.h
index 2a3164a2da..680bd8f207 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_device.h
+++ b/level_zero/core/test/unit_tests/mocks/mock_device.h
@@ -18,6 +18,8 @@
 #include "level_zero/core/test/unit_tests/white_box.h"
 #include "level_zero/tools/source/metrics/metric.h"
 
+#include "gmock/gmock.h"
+
 namespace L0 {
 namespace ult {
 
diff --git a/level_zero/core/test/unit_tests/mocks/mock_driver.cpp b/level_zero/core/test/unit_tests/mocks/mock_driver.cpp
index c114e9f68a..8db329cba2 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_driver.cpp
+++ b/level_zero/core/test/unit_tests/mocks/mock_driver.cpp
@@ -11,13 +11,10 @@ namespace L0 {
 namespace ult {
 
 using MockDriver = Mock<L0::ult::Driver>;
-using ::testing::Invoke;
 
 Mock<Driver>::Mock() {
     previousDriver = driver;
     driver = this;
-    EXPECT_CALL(*this, initialize)
-        .WillRepeatedly(Invoke(this, &MockDriver::mockInitialize));
 }
 
 Mock<Driver>::~Mock() {
diff --git a/level_zero/core/test/unit_tests/mocks/mock_driver.h b/level_zero/core/test/unit_tests/mocks/mock_driver.h
index 61a3de47fb..2a67e52cde 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_driver.h
+++ b/level_zero/core/test/unit_tests/mocks/mock_driver.h
@@ -26,19 +26,13 @@ struct Mock<Driver> : public Driver {
     Mock();
     ~Mock() override;
 
-    MOCK_METHOD(ze_result_t,
-                driverInit,
-                (ze_init_flags_t),
-                (override));
-    MOCK_METHOD(void,
-                initialize,
-                (ze_result_t * result),
-                (override));
-
-    ze_result_t mockInit(ze_init_flag_t) { return this->DriverImp::driverInit(ZE_INIT_FLAG_GPU_ONLY); }
-    void mockInitialize(ze_result_t *result) { *result = ZE_RESULT_SUCCESS; }
+    ze_result_t driverInit(ze_init_flags_t flag) override {
+        initCalledCount++;
+        return ZE_RESULT_SUCCESS;
+    }
 
     Driver *previousDriver = nullptr;
+    uint32_t initCalledCount = 0;
 };
 
 } // namespace ult
diff --git a/level_zero/core/test/unit_tests/mocks/mock_event.h b/level_zero/core/test/unit_tests/mocks/mock_event.h
index e41882aa92..98b7003325 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_event.h
+++ b/level_zero/core/test/unit_tests/mocks/mock_event.h
@@ -11,6 +11,8 @@
 #include "level_zero/core/test/unit_tests/mock.h"
 #include "level_zero/core/test/unit_tests/white_box.h"
 
+#include "gmock/gmock.h"
+
 #include <vector>
 
 #if defined(__clang__)
diff --git a/level_zero/core/test/unit_tests/mocks/mock_memory_manager.cpp b/level_zero/core/test/unit_tests/mocks/mock_memory_manager.cpp
deleted file mode 100644
index 215e9312fe..0000000000
--- a/level_zero/core/test/unit_tests/mocks/mock_memory_manager.cpp
+++ /dev/null
@@ -1,55 +0,0 @@
-/*
- * Copyright (C) 2019-2020 Intel Corporation
- *
- * SPDX-License-Identifier: MIT
- *
- */
-
-#include "mock_memory_manager.h"
-
-#include "shared/source/helpers/aligned_memory.h"
-#include "shared/source/helpers/string.h"
-
-#include <iostream>
-
-namespace L0 {
-namespace ult {
-
-using namespace testing;
-using ::testing::Return;
-
-Mock<NEO::MemoryManager>::Mock(NEO::ExecutionEnvironment &executionEnvironment) : MemoryManagerMock(executionEnvironment) {
-    EXPECT_CALL(*this, allocateGraphicsMemoryInPreferredPool)
-        .WillRepeatedly(Invoke(this, &Mock<NEO::MemoryManager>::doAllocateGraphicsMemoryInPreferredPool));
-    EXPECT_CALL(*this, freeGraphicsMemory)
-        .WillRepeatedly(Invoke(this, &Mock<NEO::MemoryManager>::doFreeGraphicsMemory));
-}
-
-inline NEO::GraphicsAllocation *Mock<NEO::MemoryManager>::doAllocateGraphicsMemoryInPreferredPool(const NEO::AllocationProperties &properties, const void *hostPtr) {
-    void *buffer;
-    if (hostPtr != nullptr) {
-        buffer = const_cast<void *>(hostPtr);
-    } else {
-        buffer = alignedMalloc(properties.size, MemoryConstants::pageSize);
-    }
-    uint64_t baseAddress = 0;
-    if (properties.allocationType == NEO::GraphicsAllocation::AllocationType::INTERNAL_HEAP) {
-        baseAddress = castToUint64(alignDown(buffer, MemoryConstants::pageSize64k));
-    }
-    auto allocation = new NEO::GraphicsAllocation(properties.rootDeviceIndex, properties.allocationType,
-                                                  buffer, reinterpret_cast<uint64_t>(buffer), baseAddress, properties.size,
-                                                  MemoryPool::System4KBPages);
-
-    return allocation;
-}
-
-inline void Mock<NEO::MemoryManager>::doFreeGraphicsMemory(NEO::GraphicsAllocation *allocation) {
-    if (allocation == nullptr) {
-        return;
-    }
-    alignedFree(allocation->getUnderlyingBuffer());
-    delete allocation;
-}
-
-} // namespace ult
-} // namespace L0
diff --git a/level_zero/core/test/unit_tests/mocks/mock_memory_manager.h b/level_zero/core/test/unit_tests/mocks/mock_memory_manager.h
index 6b2e17362a..63144704da 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_memory_manager.h
+++ b/level_zero/core/test/unit_tests/mocks/mock_memory_manager.h
@@ -32,17 +32,6 @@ struct WhiteBox<::NEO::OsAgnosticMemoryManager> : public ::NEO::OsAgnosticMemory
 
 using MemoryManagerMock = WhiteBox<::NEO::OsAgnosticMemoryManager>;
 
-template <>
-struct Mock<NEO::MemoryManager> : public MemoryManagerMock {
-    Mock(NEO::ExecutionEnvironment &executionEnvironment);
-    MOCK_METHOD2(allocateGraphicsMemoryInPreferredPool,
-                 NEO::GraphicsAllocation *(const NEO::AllocationProperties &properties, const void *hostPtr));
-    MOCK_METHOD1(freeGraphicsMemory, void(NEO::GraphicsAllocation *));
-
-    NEO::GraphicsAllocation *doAllocateGraphicsMemoryInPreferredPool(const NEO::AllocationProperties &properties, const void *hostPtr);
-    void doFreeGraphicsMemory(NEO::GraphicsAllocation *allocation);
-};
-
 } // namespace ult
 } // namespace L0
 
diff --git a/level_zero/core/test/unit_tests/mocks/mock_module.h b/level_zero/core/test/unit_tests/mocks/mock_module.h
index a4bb73644a..99e614c3bc 100644
--- a/level_zero/core/test/unit_tests/mocks/mock_module.h
+++ b/level_zero/core/test/unit_tests/mocks/mock_module.h
@@ -14,6 +14,7 @@
 #include "level_zero/core/test/unit_tests/mock.h"
 #include "level_zero/core/test/unit_tests/white_box.h"
 
+#include "gmock/gmock.h"
 namespace L0 {
 namespace ult {
 
diff --git a/level_zero/core/test/unit_tests/sources/driver/test_driver.cpp b/level_zero/core/test/unit_tests/sources/driver/test_driver.cpp
index b28cfa0c73..d6bdb045e9 100644
--- a/level_zero/core/test/unit_tests/sources/driver/test_driver.cpp
+++ b/level_zero/core/test/unit_tests/sources/driver/test_driver.cpp
@@ -15,12 +15,22 @@
 #include "level_zero/core/source/driver/driver_handle_imp.h"
 #include "level_zero/core/source/driver/driver_imp.h"
 #include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
+#include "level_zero/core/test/unit_tests/mocks/mock_driver.h"
 
 #include <bitset>
 
 namespace L0 {
 namespace ult {
 
+TEST(zeInit, whenCallingZeInitThenInitializeOnDriverIsCalled) {
+    Mock<Driver> driver;
+
+    auto result = zeInit(ZE_INIT_FLAG_GPU_ONLY);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    EXPECT_EQ(1u, driver.initCalledCount);
+}
+
 using DriverVersionTest = Test<DeviceFixture>;
 
 TEST_F(DriverVersionTest, givenCallToGetExtensionPropertiesThenUnsupportedIsReturned) {