From bd8a488f4b6d1ad3b3ce9ad0bc1e772d136a805c Mon Sep 17 00:00:00 2001
From: Edward Chen <18449977+edgchen1@users.noreply.github.com>
Date: Tue, 22 Aug 2023 12:13:52 -0700
Subject: [PATCH] Enable verbose logging in unit test program with environment
 variable. (#17133)

Enable verbose logging in unit test program with environment variable.
E.g., `ORT_UNIT_TEST_MAIN_LOG_LEVEL=0 ./onnxruntime_test_all --gtest_filter="<test that I want to see more logs for>"`.
---
 cmake/onnxruntime_unittests.cmake           |  9 ++++++++
 onnxruntime/test/unittest_main/test_main.cc | 23 ++++++++++++++++++---
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake
index 695e218441..3b9727ec08 100644
--- a/cmake/onnxruntime_unittests.cmake
+++ b/cmake/onnxruntime_unittests.cmake
@@ -1245,6 +1245,9 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
         list(APPEND onnxruntime_shared_lib_test_LIBS onnxruntime_providers_snpe)
       endif()
     endif()
+    if (CPUINFO_SUPPORTED AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
+      list(APPEND onnxruntime_shared_lib_test_LIBS cpuinfo)
+    endif()
     if (onnxruntime_USE_CUDA)
       list(APPEND onnxruntime_shared_lib_test_LIBS onnxruntime_test_cuda_ops_lib cudart)
     endif()
@@ -1543,6 +1546,12 @@ if (NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
             ${ONNXRUNTIME_CUSTOM_OP_REGISTRATION_TEST_SRC_DIR}/test_registercustomops.cc)
 
     set(onnxruntime_customopregistration_test_LIBS custom_op_library onnxruntime_common onnxruntime_test_utils)
+    if (NOT WIN32)
+      list(APPEND onnxruntime_customopregistration_test_LIBS nsync::nsync_cpp)
+    endif()
+    if (CPUINFO_SUPPORTED AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
+      list(APPEND onnxruntime_customopregistration_test_LIBS cpuinfo)
+    endif()
     if (onnxruntime_USE_TENSORRT)
       list(APPEND onnxruntime_customopregistration_test_LIBS ${TENSORRT_LIBRARY_INFER})
     endif()
diff --git a/onnxruntime/test/unittest_main/test_main.cc b/onnxruntime/test/unittest_main/test_main.cc
index b44ffc4f99..97169df36f 100644
--- a/onnxruntime/test/unittest_main/test_main.cc
+++ b/onnxruntime/test/unittest_main/test_main.cc
@@ -1,12 +1,16 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <algorithm>
+#include <cstdlib>
+#include <optional>
+#include <string>
+
 #ifndef USE_ONNXRUNTIME_DLL
 #ifdef __GNUC__
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wignored-qualifiers"
 #pragma GCC diagnostic ignored "-Wunused-parameter"
-#else
 #endif
 #include <google/protobuf/message_lite.h>
 #ifdef __GNUC__
@@ -14,9 +18,11 @@
 #endif
 #endif
 
+#include "gtest/gtest.h"
+
+#include "core/platform/env_var_utils.h"
 #include "core/session/onnxruntime_cxx_api.h"
 #include "core/util/thread_utils.h"
-#include "gtest/gtest.h"
 #include "test/test_environment.h"
 
 std::unique_ptr<Ort::Env> ort_env;
@@ -54,8 +60,19 @@ int TEST_MAIN(int argc, char** argv) {
 
   ORT_TRY {
     ::testing::InitGoogleTest(&argc, argv);
-
     ortenv_setup();
+
+    // allow verbose logging to be enabled by setting this environment variable to a numeric log level
+    constexpr auto kLogLevelEnvironmentVariableName = "ORT_UNIT_TEST_MAIN_LOG_LEVEL";
+    if (auto log_level = onnxruntime::ParseEnvironmentVariable<int>(kLogLevelEnvironmentVariableName);
+        log_level.has_value()) {
+      *log_level = std::clamp(*log_level,
+                              static_cast<int>(ORT_LOGGING_LEVEL_VERBOSE),
+                              static_cast<int>(ORT_LOGGING_LEVEL_FATAL));
+      std::cout << "Setting log level to " << *log_level << "\n";
+      ort_env->UpdateEnvWithCustomLogLevel(static_cast<OrtLoggingLevel>(*log_level));
+    }
+
     status = RUN_ALL_TESTS();
   }
   ORT_CATCH(const std::exception& ex) {