Made TestSink prettier

This commit is contained in:
Aleksey Vitebskiy 2020-11-02 14:30:39 -05:00
parent 190b8f380a
commit bbb3cb142e

View File

@ -99,12 +99,18 @@ TEST(InternalLogString, Basic) {
internal::LogString(__FILE__, __LINE__, INFO, "Hello there"); internal::LogString(__FILE__, __LINE__, INFO, "Hello there");
} }
struct TestSink : public TFLogSink { class TestSink : public TFLogSink {
std::stringstream ss; public:
void Send(const TFLogEntry& entry) override { void Send(const TFLogEntry& entry) override {
ss << entry.ToString() << std::endl; ss << entry.ToString() << std::endl;
} }
std::string Get() const {
return ss.str();
}
private:
std::stringstream ss;
}; };
} // namespace tensorflow } // namespace tensorflow
@ -213,10 +219,10 @@ TEST_F(LogSinkTest, testLogSinks) {
sinks = TFGetLogSinks(); sinks = TFGetLogSinks();
ASSERT_EQ(sinks.size(), 1); ASSERT_EQ(sinks.size(), 1);
ASSERT_EQ(sink.ss.str(), "Foo\n"); ASSERT_EQ(sink.Get(), "Foo\n");
LOG(INFO) << "Bar"; LOG(INFO) << "Bar";
ASSERT_EQ(sink.ss.str(), "Foo\nBar\n"); ASSERT_EQ(sink.Get(), "Foo\nBar\n");
} }
} // namespace tensorflow } // namespace tensorflow
@ -247,7 +253,7 @@ TEST(LogSinkTest, testLogSinks)
LOG(INFO) << "Foo"; LOG(INFO) << "Foo";
LOG(INFO) << "Bar"; LOG(INFO) << "Bar";
EXPECT_EQ(sink.ss.str(), "Foo\nBar\n"); EXPECT_EQ(sink.Get(), "Foo\nBar\n");
TFRemoveLogSink(&sink); TFRemoveLogSink(&sink);
sinks = TFGetLogSinks(); sinks = TFGetLogSinks();