| // Copyright 2023 The Centipede Authors. |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // https://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| #include "./centipede/distill.h" |
| |
| #include <filesystem> // NOLINT |
| #include <vector> |
| |
| #include "gmock/gmock.h" |
| #include "gtest/gtest.h" |
| #include "absl/flags/flag.h" |
| #include "absl/flags/reflection.h" |
| #include "./centipede/blob_file.h" |
| #include "./centipede/defs.h" |
| #include "./centipede/environment.h" |
| #include "./centipede/feature.h" |
| #include "./centipede/test_util.h" |
| |
| ABSL_DECLARE_FLAG(std::string, binary_hash); |
| ABSL_DECLARE_FLAG(std::string, binary); |
| ABSL_DECLARE_FLAG(std::string, workdir); |
| |
| namespace centipede { |
| namespace { |
| |
| struct TestCorpusRecord { |
| ByteArray input; |
| FeatureVec feature_vec; |
| }; |
| |
| using Shard = std::vector<TestCorpusRecord>; |
| using ShardVec = std::vector<Shard>; |
| using InputVec = std::vector<ByteArray>; |
| |
| // Writes `record` to shard `shard_index`. |
| void WriteToShard(const Environment &env, const TestCorpusRecord &record, |
| size_t shard_index) { |
| auto corpus_path = env.MakeCorpusPath(shard_index); |
| auto features_path = env.MakeFeaturesPath(shard_index); |
| const auto corpus_appender = DefaultBlobFileWriterFactory(); |
| const auto features_appender = DefaultBlobFileWriterFactory(); |
| CHECK_OK(corpus_appender->Open(corpus_path, "a")); |
| CHECK_OK(features_appender->Open(features_path, "a")); |
| CHECK_OK(corpus_appender->Write(record.input)); |
| CHECK_OK(features_appender->Write( |
| PackFeaturesAndHash(record.input, record.feature_vec))); |
| } |
| |
| // Reads and returns the distilled corpus from `env.MakeDistilledPath()`. |
| std::vector<ByteArray> ReadFromDistilled(const Environment &env) { |
| auto distilled_path = env.MakeDistilledPath(); |
| auto reader = DefaultBlobFileReaderFactory(); |
| CHECK_OK(reader->Open(distilled_path)); |
| absl::Span<uint8_t> blob; |
| std::vector<ByteArray> result; |
| while (reader->Read(blob).ok()) { |
| result.emplace_back(blob.begin(), blob.end()); |
| } |
| return result; |
| } |
| |
| // Distills `shards` in the order specified by `shard_indices`, |
| // returns the distilled corpus as a vector of inputs. |
| InputVec TestDistill(const ShardVec &shards, |
| const std::vector<size_t> &shard_indices, |
| std::string_view test_name) { |
| // Set up the environment. |
| // We need to set at least --binary_hash before `env` is constructed, |
| // so we do this by overriding the flags. |
| absl::FlagSaver flag_saver; |
| std::string dir = std::filesystem::path(GetTestTempDir()).append(test_name); |
| std::filesystem::remove_all(dir); |
| std::filesystem::create_directories(dir); |
| absl::SetFlag(&FLAGS_workdir, dir); |
| absl::SetFlag(&FLAGS_binary, "binary_that_is_not_here"); |
| absl::SetFlag(&FLAGS_binary_hash, "01234567890"); |
| Environment env; |
| env.total_shards = shards.size(); |
| env.my_shard_index = 1; // an arbitrary shard index. |
| std::filesystem::create_directories(env.MakeCoverageDirPath()); |
| |
| // Write the shards. |
| for (size_t shard_index = 0; shard_index < shards.size(); ++shard_index) { |
| for (const auto &record : shards[shard_index]) { |
| WriteToShard(env, record, shard_index); |
| } |
| } |
| // Distill. |
| DistillTask(env, shard_indices); |
| // Read the result back. |
| return ReadFromDistilled(env); |
| } |
| |
| TEST(Distill, BasicDistill) { |
| ByteArray in0 = {0}; |
| ByteArray in1 = {1}; |
| ByteArray in2 = {2}; |
| ByteArray in3 = {3}; |
| ShardVec shards = { |
| // shard 0; note: distillation iterates the shards backwards. |
| {{in3, {10}}, {in0, {10, 20}}}, |
| // shard 1 |
| {{in1, {20, 30}}}, |
| // shard 2 |
| {{in2, {30, 40}}}, |
| }; |
| // Distill these 3 shards in different orders, observe different results. |
| EXPECT_THAT(TestDistill(shards, {0, 1, 2}, test_info_->name()), |
| testing::ElementsAreArray({in0, in1, in2})); |
| EXPECT_THAT(TestDistill(shards, {2, 0, 1}, test_info_->name()), |
| testing::ElementsAreArray({in2, in0})); |
| EXPECT_THAT(TestDistill(shards, {1, 0, 2}, test_info_->name()), |
| testing::ElementsAreArray({in1, in0, in2})); |
| } |
| |
| // TODO(kcc): add more tests once we settle on the testing code above. |
| |
| } // namespace |
| } // namespace centipede |