diff --git a/thinkit/BUILD.bazel b/thinkit/BUILD.bazel index 4a52be59..3e0935fe 100644 --- a/thinkit/BUILD.bazel +++ b/thinkit/BUILD.bazel @@ -188,14 +188,12 @@ cc_library( hdrs = ["bazel_test_environment.h"], deps = [ ":test_environment", - "//gutil:status", + "//gutil:test_artifact_writer", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", - "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", - "@com_google_googletest//:gtest", ], ) diff --git a/thinkit/bazel_test_environment.cc b/thinkit/bazel_test_environment.cc index 2b5fb936..ba7d7402 100644 --- a/thinkit/bazel_test_environment.cc +++ b/thinkit/bazel_test_environment.cc @@ -14,106 +14,31 @@ #include "thinkit/bazel_test_environment.h" -#include -#include -#include #include -#include -#include -#include // NOLINT -#include -#include "absl/container/flat_hash_map.h" #include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" -#include "gtest/gtest.h" -#include "gutil/status.h" -#include "thinkit/test_environment.h" +#include "gutil/test_artifact_writer.h" namespace thinkit { -namespace { - -absl::StatusOr ArtifactDirectory() { - // Pick appropriate artifact directory using Bazel environment variables, see - // https://docs.bazel.build/versions/main/test-encyclopedia.html#initial-conditions - char* base_dir = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); - if (base_dir == nullptr) { - base_dir = std::getenv("TEST_TMPDIR"); - } - if (base_dir == nullptr) { - return gutil::InternalErrorBuilder() - << "Environment variables TEST_UNDECLARED_OUTPUTS_DIR and " - "TEST_TMPDIR undefined; is this a Bazel test run?"; - } - std::string dir = base_dir; - if (auto* test_info = testing::UnitTest::GetInstance()->current_test_info(); - test_info != nullptr) { - absl::StrAppend(&dir, "/", test_info->test_case_name(), "/", - test_info->name()); - } - - // Ensure the directory exists. - std::error_code error; - std::filesystem::create_directories(dir, error); - if (error) { - return gutil::InternalErrorBuilder() - << "failed to create test artifact directory '" << dir - << "': " << error; - } - return dir; -} -absl::Status WriteToTestArtifact( - absl::string_view filename, absl::string_view contents, - std::ios_base::openmode mode, - absl::flat_hash_map& open_file_by_filepath) { - ASSIGN_OR_RETURN(std::string directory, ArtifactDirectory()); - std::string filepath = absl::StrCat(directory, "/", filename); - // Note that pointer-stability of values is not a concern here because the - // reference is local and nothing is added to the map while the reference is - // live. - std::ofstream& file = open_file_by_filepath[filepath]; - if (file.is_open() && mode == std::ios_base::trunc) { - // If we have an open file descriptor and we want to truncate the file, then - // we close it and reopen it in truncation mode. - file.close(); - } - // If the file is not open, then we just want to open it in the given mode. - if (!file.is_open()) { - file.open(filepath, mode); - // If the file is still not open, we have a problem. - if (!file.is_open()) { - return gutil::InternalErrorBuilder() - << "unable to open test artifact file: '" << filepath << "'"; - } - } - file << contents; - - // We flush the contents to persist them, but leave the file open so that we - // can continue to append to it. - file.flush(); - if (file.good()) return absl::OkStatus(); - return gutil::InternalErrorBuilder() - << "failed to store test artifact: '" << filepath << "'"; -} - -} // namespace - absl::Status BazelTestEnvironment::StoreTestArtifact( absl::string_view filename, absl::string_view contents) { - absl::MutexLock lock(&this->write_mutex_); - return WriteToTestArtifact(filename, contents, std::ios_base::trunc, - open_file_by_filepath_); + return artifact_writer_.StoreTestArtifact(filename, contents); +} +absl::Status BazelTestEnvironment::StoreTestArtifact( + absl::string_view filename, const google::protobuf::Message& proto) { + return artifact_writer_.StoreTestArtifact(filename, proto); } absl::Status BazelTestEnvironment::AppendToTestArtifact( absl::string_view filename, absl::string_view contents) { - absl::MutexLock lock(&this->write_mutex_); - return WriteToTestArtifact(filename, contents, std::ios_base::app, - open_file_by_filepath_); + return artifact_writer_.AppendToTestArtifact(filename, contents); +} +absl::Status BazelTestEnvironment::AppendToTestArtifact( + absl::string_view filename, const google::protobuf::Message& proto) { + return artifact_writer_.AppendToTestArtifact(filename, proto); } } // namespace thinkit diff --git a/thinkit/bazel_test_environment.h b/thinkit/bazel_test_environment.h index 19af59b0..d7df1ce8 100644 --- a/thinkit/bazel_test_environment.h +++ b/thinkit/bazel_test_environment.h @@ -27,6 +27,7 @@ #include "absl/status/status.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" +#include "gutil/test_artifact_writer.h" #include "thinkit/test_environment.h" namespace thinkit { @@ -46,17 +47,14 @@ class BazelTestEnvironment : public TestEnvironment { set_test_case_ids_(std::move(set_test_case_ids)) {} absl::Status StoreTestArtifact(absl::string_view filename, - absl::string_view contents) - ABSL_LOCKS_EXCLUDED(write_mutex_) override; - using TestEnvironment::StoreTestArtifact; // Inherit protobuf overload which - // calls string_view version. + absl::string_view contents) override; + absl::Status StoreTestArtifact(absl::string_view filename, + const google::protobuf::Message& proto); absl::Status AppendToTestArtifact(absl::string_view filename, - absl::string_view contents) - ABSL_LOCKS_EXCLUDED(write_mutex_) override; - using TestEnvironment::AppendToTestArtifact; // Inherit protobuf overload - // which calls string_view - // version. + absl::string_view contents) override; + absl::Status AppendToTestArtifact(absl::string_view filename, + const google::protobuf::Message& proto); bool MaskKnownFailures() { return mask_known_failures_; }; @@ -67,12 +65,7 @@ class BazelTestEnvironment : public TestEnvironment { private: bool mask_known_failures_; std::function&)> set_test_case_ids_; - // Open files are cached to avoid closing them after every append. On certain - // file systems (e.g. b/111316875) closing files is abnormally slow and this - // avoids it. However, this approach should also generally be faster. - absl::flat_hash_map open_file_by_filepath_; - // The mutex is used to ensure that writes to disk are sequential. - absl::Mutex write_mutex_; + gutil::BazelTestArtifactWriter artifact_writer_; }; } // namespace thinkit diff --git a/thinkit/bazel_test_environment_test.cc b/thinkit/bazel_test_environment_test.cc index 55d9ffba..e45ed856 100644 --- a/thinkit/bazel_test_environment_test.cc +++ b/thinkit/bazel_test_environment_test.cc @@ -106,7 +106,7 @@ TEST(BazelTestEnvironmentTest, SetTestCaseIdWorksForUnaryConstructor) { // -- Benchmarks --------------------------------------------------------------- // -// Best run with 'blaze test --benchmarks=all'. +// Best run with 'blaze test --test_arg=--benchmark_filter=all '. // // Ideally, we would like to use 'benchy' for benchmarking purposes, but, // because we are benchmarking a testing environment, it relies on being set