• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 #include <gtest/gtest.h>
2 #include <test/cpp/jit/test_utils.h>
3 
4 #include <torch/csrc/jit/jit_log.h>
5 #include <sstream>
6 
7 namespace torch {
8 namespace jit {
9 
TEST(JitLoggingTest,CheckSetLoggingLevel)10 TEST(JitLoggingTest, CheckSetLoggingLevel) {
11   ::torch::jit::set_jit_logging_levels("file_to_test");
12   ASSERT_TRUE(::torch::jit::is_enabled(
13       "file_to_test.cpp", JitLoggingLevels::GRAPH_DUMP));
14 }
15 
TEST(JitLoggingTest,CheckSetMultipleLogLevels)16 TEST(JitLoggingTest, CheckSetMultipleLogLevels) {
17   ::torch::jit::set_jit_logging_levels("f1:>f2:>>f3");
18   ASSERT_TRUE(::torch::jit::is_enabled("f1.cpp", JitLoggingLevels::GRAPH_DUMP));
19   ASSERT_TRUE(
20       ::torch::jit::is_enabled("f2.cpp", JitLoggingLevels::GRAPH_UPDATE));
21   ASSERT_TRUE(
22       ::torch::jit::is_enabled("f3.cpp", JitLoggingLevels::GRAPH_DEBUG));
23 }
24 
TEST(JitLoggingTest,CheckLoggingLevelAfterUnset)25 TEST(JitLoggingTest, CheckLoggingLevelAfterUnset) {
26   ::torch::jit::set_jit_logging_levels("f1");
27   ASSERT_EQ("f1", ::torch::jit::get_jit_logging_levels());
28   ::torch::jit::set_jit_logging_levels("invalid");
29   ASSERT_FALSE(
30       ::torch::jit::is_enabled("f1.cpp", JitLoggingLevels::GRAPH_DUMP));
31 }
32 
TEST(JitLoggingTest,CheckAfterChangingLevel)33 TEST(JitLoggingTest, CheckAfterChangingLevel) {
34   ::torch::jit::set_jit_logging_levels("f1");
35   ::torch::jit::set_jit_logging_levels(">f1");
36   ASSERT_TRUE(
37       ::torch::jit::is_enabled("f1.cpp", JitLoggingLevels::GRAPH_UPDATE));
38 }
39 
TEST(JitLoggingTest,CheckOutputStreamSetting)40 TEST(JitLoggingTest, CheckOutputStreamSetting) {
41   ::torch::jit::set_jit_logging_levels("test_jit_logging_levels");
42   std::ostringstream test_stream;
43   ::torch::jit::set_jit_logging_output_stream(test_stream);
44   /* Using JIT_LOG checks if this file has logging enabled with
45     is_enabled(__FILE__, level) making the test fail. since we are only testing
46     the OutputStreamSetting we can forcefully output to it directly.
47   */
48   ::torch::jit::get_jit_logging_output_stream() << ::torch::jit::jit_log_prefix(
49       ::torch::jit::JitLoggingLevels::GRAPH_DUMP,
50       __FILE__,
51       __LINE__,
52       ::c10::str("Message"));
53   ASSERT_TRUE(test_stream.str().size() > 0);
54 }
55 
56 } // namespace jit
57 } // namespace torch
58