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 | |
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 | |
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 | |
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 | |
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 | |
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 | |