From 8daf71070d680d371b94c4eb71a4c7f7b0d878d4 Mon Sep 17 00:00:00 2001
From: Mateusz Hoppe <mateusz.hoppe@intel.com>
Date: Thu, 30 Jul 2020 13:03:04 +0200
Subject: [PATCH] Enhance mocking env variables

Change-Id: I2d0659b6f608467f9c1bc5746742ac5e3454582d
Signed-off-by: Mateusz Hoppe <mateusz.hoppe@intel.com>
---
 .../test/unit_tests/sources/driver/test_driver.cpp   | 11 ++++-------
 opencl/test/unit_test/libult/io_functions.cpp        |  4 +---
 opencl/test/unit_test/mocks/mock_io_functions.h      | 12 ++++--------
 3 files changed, 9 insertions(+), 18 deletions(-)

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 e6768e986d..54af980ad9 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
@@ -93,9 +93,7 @@ TEST(DriverImpTest, givenDriverImpWhenInitializedThenEnvVariablesAreRead) {
     NEO::HardwareInfo hwInfo = *NEO::defaultHwInfo.get();
     hwInfo.capabilityTable.levelZeroSupported = true;
 
-    VariableBackup<bool> mockDeviceFlagBackup(&IoFunctions::returnMockEnvValue, true);
     VariableBackup<uint32_t> mockGetenvCalledBackup(&IoFunctions::mockGetenvCalled, 0);
-    VariableBackup<std::set<std::string>> notMockableEnvValuesBackup(&IoFunctions::notMockableEnvValues, {"ZE_ENABLE_METRICS"});
 
     ze_result_t result = ZE_RESULT_ERROR_UNINITIALIZED;
     DriverImp driverImp;
@@ -112,8 +110,9 @@ TEST(DriverImpTest, givenMissingMetricApiDependenciesWhenInitializingDriverImpTh
     NEO::HardwareInfo hwInfo = *NEO::defaultHwInfo.get();
     hwInfo.capabilityTable.levelZeroSupported = true;
 
-    VariableBackup<bool> mockDeviceFlagBackup(&IoFunctions::returnMockEnvValue, true);
     VariableBackup<uint32_t> mockGetenvCalledBackup(&IoFunctions::mockGetenvCalled, 0);
+    std::unordered_map<std::string, std::string> mockableEnvs = {{"ZE_ENABLE_METRICS", "1"}};
+    VariableBackup<std::unordered_map<std::string, std::string> *> mockableEnvValuesBackup(&IoFunctions::mockableEnvValues, &mockableEnvs);
 
     ze_result_t result = ZE_RESULT_ERROR_UNINITIALIZED;
     DriverImp driverImp;
@@ -127,9 +126,9 @@ TEST(DriverImpTest, givenEnabledProgramDebuggingWhenCreatingExecutionEnvironment
     NEO::HardwareInfo hwInfo = *NEO::defaultHwInfo.get();
     hwInfo.capabilityTable.levelZeroSupported = true;
 
-    VariableBackup<bool> mockDeviceFlagBackup(&IoFunctions::returnMockEnvValue, true);
     VariableBackup<uint32_t> mockGetenvCalledBackup(&IoFunctions::mockGetenvCalled, 0);
-    VariableBackup<std::set<std::string>> notMockableEnvValuesBackup(&IoFunctions::notMockableEnvValues, {"ZE_ENABLE_METRICS"});
+    std::unordered_map<std::string, std::string> mockableEnvs = {{"ZET_ENABLE_PROGRAM_DEBUGGING", "1"}};
+    VariableBackup<std::unordered_map<std::string, std::string> *> mockableEnvValuesBackup(&IoFunctions::mockableEnvValues, &mockableEnvs);
 
     ze_result_t result = ZE_RESULT_ERROR_UNINITIALIZED;
     DriverImp driverImp;
@@ -148,8 +147,6 @@ TEST(DriverImpTest, givenNoProgramDebuggingEnvVarWhenCreatingExecutionEnvironmen
     NEO::HardwareInfo hwInfo = *NEO::defaultHwInfo.get();
     hwInfo.capabilityTable.levelZeroSupported = true;
 
-    VariableBackup<bool> mockDeviceFlagBackup(&IoFunctions::returnMockEnvValue, false);
-
     ze_result_t result = ZE_RESULT_ERROR_UNINITIALIZED;
     DriverImp driverImp;
     driverImp.initialize(&result);
diff --git a/opencl/test/unit_test/libult/io_functions.cpp b/opencl/test/unit_test/libult/io_functions.cpp
index 2f39e262ef..1ed13cc4a7 100644
--- a/opencl/test/unit_test/libult/io_functions.cpp
+++ b/opencl/test/unit_test/libult/io_functions.cpp
@@ -19,9 +19,7 @@ uint32_t mockVfptrinfCalled = 0;
 uint32_t mockFcloseCalled = 0;
 uint32_t mockGetenvCalled = 0;
 
-bool returnMockEnvValue = false;
-std::string mockEnvValue = "1";
-std::set<std::string> notMockableEnvValues = {""};
+std::unordered_map<std::string, std::string> *mockableEnvValues = nullptr;
 
 } // namespace IoFunctions
 } // namespace NEO
diff --git a/opencl/test/unit_test/mocks/mock_io_functions.h b/opencl/test/unit_test/mocks/mock_io_functions.h
index 490b7eb134..66393f5cd4 100644
--- a/opencl/test/unit_test/mocks/mock_io_functions.h
+++ b/opencl/test/unit_test/mocks/mock_io_functions.h
@@ -9,8 +9,8 @@
 #include "shared/source/utilities/io_functions.h"
 
 #include <cstdint>
-#include <set>
 #include <string>
+#include <unordered_map>
 
 namespace NEO {
 namespace IoFunctions {
@@ -19,9 +19,7 @@ extern uint32_t mockVfptrinfCalled;
 extern uint32_t mockFcloseCalled;
 extern uint32_t mockGetenvCalled;
 
-extern bool returnMockEnvValue;
-extern std::string mockEnvValue;
-extern std::set<std::string> notMockableEnvValues;
+extern std::unordered_map<std::string, std::string> *mockableEnvValues;
 
 inline FILE *mockFopen(const char *filename, const char *mode) {
     mockFopenCalled++;
@@ -40,10 +38,8 @@ inline int mockFclose(FILE *stream) {
 
 inline char *mockGetenv(const char *name) noexcept {
     mockGetenvCalled++;
-    if (notMockableEnvValues.find(name) == notMockableEnvValues.end()) {
-        if (returnMockEnvValue) {
-            return const_cast<char *>(mockEnvValue.c_str());
-        }
+    if (mockableEnvValues != nullptr && mockableEnvValues->find(name) != mockableEnvValues->end()) {
+        return const_cast<char *>(mockableEnvValues->find(name)->second.c_str());
     }
     return getenv(name);
 }