1 #include <gmock/gmock.h>
2 #include <gtest/gtest.h>
3
4 #include <ATen/ATen.h>
5 #include <ATen/core/Vitals.h>
6 #include <c10/util/irange.h>
7 #include <cstdlib>
8
9 using namespace at::vitals;
10 using ::testing::HasSubstr;
11
TEST(Vitals,Basic)12 TEST(Vitals, Basic) {
13 std::stringstream buffer;
14
15 std::streambuf* sbuf = std::cout.rdbuf();
16 std::cout.rdbuf(buffer.rdbuf());
17 {
18 #ifdef _WIN32
19 _putenv("TORCH_VITAL=1");
20 #else
21 setenv("TORCH_VITAL", "1", 1);
22 #endif
23 TORCH_VITAL_DEFINE(Testing);
24 TORCH_VITAL(Testing, Attribute0) << 1;
25 TORCH_VITAL(Testing, Attribute1) << "1";
26 TORCH_VITAL(Testing, Attribute2) << 1.0f;
27 TORCH_VITAL(Testing, Attribute3) << 1.0;
28 auto t = at::ones({1, 1});
29 TORCH_VITAL(Testing, Attribute4) << t;
30 }
31 std::cout.rdbuf(sbuf);
32
33 auto s = buffer.str();
34 ASSERT_THAT(s, HasSubstr("Testing.Attribute0\t\t 1"));
35 ASSERT_THAT(s, HasSubstr("Testing.Attribute1\t\t 1"));
36 ASSERT_THAT(s, HasSubstr("Testing.Attribute2\t\t 1"));
37 ASSERT_THAT(s, HasSubstr("Testing.Attribute3\t\t 1"));
38 ASSERT_THAT(s, HasSubstr("Testing.Attribute4\t\t 1"));
39 }
40
TEST(Vitals,MultiString)41 TEST(Vitals, MultiString) {
42 std::stringstream buffer;
43
44 std::streambuf* sbuf = std::cout.rdbuf();
45 std::cout.rdbuf(buffer.rdbuf());
46 {
47 #ifdef _WIN32
48 _putenv("TORCH_VITAL=1");
49 #else
50 setenv("TORCH_VITAL", "1", 1);
51 #endif
52 TORCH_VITAL_DEFINE(Testing);
53 TORCH_VITAL(Testing, Attribute0) << 1 << " of " << 2;
54 TORCH_VITAL(Testing, Attribute1) << 1;
55 TORCH_VITAL(Testing, Attribute1) << " of ";
56 TORCH_VITAL(Testing, Attribute1) << 2;
57 }
58 std::cout.rdbuf(sbuf);
59
60 auto s = buffer.str();
61 ASSERT_THAT(s, HasSubstr("Testing.Attribute0\t\t 1 of 2"));
62 ASSERT_THAT(s, HasSubstr("Testing.Attribute1\t\t 1 of 2"));
63 }
64
TEST(Vitals,OnAndOff)65 TEST(Vitals, OnAndOff) {
66 for (const auto i : c10::irange(2)) {
67 std::stringstream buffer;
68
69 std::streambuf* sbuf = std::cout.rdbuf();
70 std::cout.rdbuf(buffer.rdbuf());
71 {
72 #ifdef _WIN32
73 if (i) {
74 _putenv("TORCH_VITAL=1");
75 } else {
76 _putenv("TORCH_VITAL=0");
77 }
78 #else
79 setenv("TORCH_VITAL", i ? "1" : "", 1);
80 #endif
81 TORCH_VITAL_DEFINE(Testing);
82 TORCH_VITAL(Testing, Attribute0) << 1;
83 }
84 std::cout.rdbuf(sbuf);
85
86 auto s = buffer.str();
87 auto f = s.find("Testing.Attribute0\t\t 1");
88 if (i) {
89 ASSERT_TRUE(f != std::string::npos);
90 } else {
91 ASSERT_TRUE(f == std::string::npos);
92 }
93 }
94 }
95
TEST(Vitals,APIVitals)96 TEST(Vitals, APIVitals) {
97 std::stringstream buffer;
98 // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
99 bool rvalue;
100 std::streambuf* sbuf = std::cout.rdbuf();
101 std::cout.rdbuf(buffer.rdbuf());
102 {
103 #ifdef _WIN32
104 _putenv("TORCH_VITAL=1");
105 #else
106 setenv("TORCH_VITAL", "1", 1);
107 #endif
108 APIVitals api_vitals;
109 rvalue = api_vitals.setVital("TestingSetVital", "TestAttr", "TestValue");
110 }
111 std::cout.rdbuf(sbuf);
112
113 auto s = buffer.str();
114 ASSERT_TRUE(rvalue);
115 ASSERT_THAT(s, HasSubstr("TestingSetVital.TestAttr\t\t TestValue"));
116 }
117