Separate out the portable log file part from the TestReporter class so that it could be used to output in other proto format.

PiperOrigin-RevId: 262411723

tensorflow/core/util/reporter.cc

@@ -21,25 +21,57 @@ limitations under the License.
 
 namespace tensorflow {
 
-TestReporter::TestReporter(const string& fname, const string& test_name)
+TestReportFile::TestReportFile(const string& fname, const string& test_name)
     : closed_(true), fname_(fname), test_name_(test_name) {}
 
-Status TestReporter::Close() {
+Status TestReportFile::Append(const string& content) {
   if (closed_) return Status::OK();
+  return log_file_->Append(content);
+}
+
+Status TestReportFile::Close() {
+  if (closed_) return Status::OK();
+  closed_ = true;
+  return log_file_->Close();
+}
+
+Status TestReportFile::Initialize() {
+  if (fname_.empty()) {
+    return Status::OK();
+  }
+  string mangled_fname = strings::StrCat(
+      fname_, absl::StrJoin(str_util::Split(test_name_, '/'), "__"));
+  Env* env = Env::Default();
+  if (env->FileExists(mangled_fname).ok()) {
+    return errors::InvalidArgument(
+        "Cannot create TestReportFile, file exists: ", mangled_fname);
+  }
+  TF_RETURN_IF_ERROR(env->NewWritableFile(mangled_fname, &log_file_));
+  TF_RETURN_IF_ERROR(log_file_->Flush());
+
+  closed_ = false;
+  return Status::OK();
+}
+
+TestReporter::TestReporter(const string& fname, const string& test_name)
+    : report_file_(fname, test_name) {
+  benchmark_entry_.set_name(test_name);
+}
+
+Status TestReporter::Close() {
+  if (report_file_.IsClosed()) return Status::OK();
 
   BenchmarkEntries entries;
   *entries.add_entry() = benchmark_entry_;
-  TF_RETURN_IF_ERROR(log_file_->Append(entries.SerializeAsString()));
+  TF_RETURN_IF_ERROR(report_file_.Append(entries.SerializeAsString()));
-
   benchmark_entry_.Clear();
-  closed_ = true;
 
-  return log_file_->Close();
+  return report_file_.Close();
 }
 
 Status TestReporter::Benchmark(int64 iters, double cpu_time, double wall_time,
                                double throughput) {
-  if (closed_) return Status::OK();
+  if (report_file_.IsClosed()) return Status::OK();
   benchmark_entry_.set_iters(iters);
   benchmark_entry_.set_cpu_time(cpu_time / iters);
   benchmark_entry_.set_wall_time(wall_time / iters);
@@ -48,34 +80,17 @@ Status TestReporter::Benchmark(int64 iters, double cpu_time, double wall_time,
 }
 
 Status TestReporter::SetProperty(const string& name, const string& value) {
-  if (closed_) return Status::OK();
+  if (report_file_.IsClosed()) return Status::OK();
   (*benchmark_entry_.mutable_extras())[name].set_string_value(value);
   return Status::OK();
 }
 
 Status TestReporter::SetProperty(const string& name, double value) {
-  if (closed_) return Status::OK();
+  if (report_file_.IsClosed()) return Status::OK();
   (*benchmark_entry_.mutable_extras())[name].set_double_value(value);
   return Status::OK();
 }
 
-Status TestReporter::Initialize() {
+Status TestReporter::Initialize() { return report_file_.Initialize(); }
-  if (fname_.empty()) {
-    return Status::OK();
-  }
-  string mangled_fname = strings::StrCat(
-      fname_, absl::StrJoin(str_util::Split(test_name_, '/'), "__"));
-  Env* env = Env::Default();
-  if (env->FileExists(mangled_fname).ok()) {
-    return errors::InvalidArgument("Cannot create TestReporter, file exists: ",
-                                   mangled_fname);
-  }
-  TF_RETURN_IF_ERROR(env->NewWritableFile(mangled_fname, &log_file_));
-  TF_RETURN_IF_ERROR(log_file_->Flush());
-
-  benchmark_entry_.set_name(test_name_);
-  closed_ = false;
-  return Status::OK();
-}
 
 }  // namespace tensorflow

tensorflow/core/util/reporter.h

@@ -29,6 +29,34 @@ limitations under the License.
 
 namespace tensorflow {
 
+// The TestReportFile provides a file abstraction for TF tests to use.
+class TestReportFile {
+ public:
+  // Create a TestReportFile with the test name 'test_name'.
+  TestReportFile(const string& fname, const string& test_name);
+
+  // Initialize the TestReportFile.  If the reporting env flag is set,
+  // try to create the reporting file.  Fails if the file already exists.
+  Status Initialize();
+
+  // Append the report file w/ 'content'.
+  Status Append(const string& content);
+
+  // Close the report file.
+  Status Close();
+
+  bool IsClosed() const { return closed_; }
+
+  ~TestReportFile() { Close().IgnoreError(); }  // Autoclose in destructor.
+
+ private:
+  bool closed_;
+  string fname_;
+  string test_name_;
+  std::unique_ptr<WritableFile> log_file_;
+  TF_DISALLOW_COPY_AND_ASSIGN(TestReportFile);
+};
+
 // The TestReporter writes test / benchmark output to binary Protobuf files when
 // the environment variable "TEST_REPORT_FILE_PREFIX" is defined.
 //
@@ -91,10 +119,7 @@ class TestReporter {
     const char* fname_ptr = getenv(kTestReporterEnv);
     return (fname_ptr != nullptr) ? fname_ptr : "";
   }
-  bool closed_;
+  TestReportFile report_file_;
-  string fname_;
-  string test_name_;
-  std::unique_ptr<WritableFile> log_file_;
   BenchmarkEntry benchmark_entry_;
   TF_DISALLOW_COPY_AND_ASSIGN(TestReporter);
 };