src/local_aggregation_1_1/local_aggregation_test.cc

// Copyright 2020 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "src/local_aggregation_1_1/local_aggregation.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "absl/strings/escaping.h"
#include "absl/strings/str_cat.h"
#include "src/lib/util/clock.h"
#include "src/lib/util/testing/test_with_files.h"
#include "src/local_aggregation_1_1/aggregation_procedures/aggregation_procedure.h"
#include "src/local_aggregation_1_1/testing/test_registry.cb.h"
#include "src/logger/event_record.h"
#include "src/logger/logger_test_utils.h"
#include "src/logger/observation_writer.h"
#include "src/logger/project_context_factory.h"
#include "src/observation_store/observation_store_internal.pb.h"
#include "src/pb/metadata_builder.h"
#include "src/registry/project.pb.h"
#include "src/system_data/client_secret.h"
#include "src/system_data/fake_system_data.h"

namespace cobalt::local_aggregation {
using ::testing::Contains;
using ::testing::Not;
using TimeInfo = util::TimeInfo;

namespace {

std::unique_ptr<CobaltRegistry> GetRegistry() {
  std::string bytes;
  if (!absl::Base64Unescape(kCobaltRegistryBase64, &bytes)) {
    LOG(ERROR) << "Unable to decode Base64 String";
    return nullptr;
  }

  auto registry = std::make_unique<CobaltRegistry>();
  if (!registry->ParseFromString(bytes)) {
    LOG(ERROR) << "Unable to parse registry from bytes";
    return nullptr;
  }

  return registry;
}

}  // namespace

class LocalAggregationTest : public util::testing::TestWithFiles {
 private:
  void SetUp() override {
    TestWithFiles::SetUp();
    global_project_context_factory_ =
        std::make_unique<logger::ProjectContextFactory>(GetRegistry());
    metadata_builder_ =
        std::make_unique<MetadataBuilder>(&system_data_, system_data_cache_path(), fs());
  }

 protected:
  std::unique_ptr<LocalAggregation> MakeLocalAggregation(
      const logger::ObservationWriter* observation_writer) {
    CobaltConfig cfg = {.local_aggregate_store_dir = local_aggregation_store_path(),
                        .client_secret = system_data::ClientSecret::GenerateNewSecret()};

    return std::make_unique<LocalAggregation>(cfg, global_project_context_factory_.get(),
                                              metadata_builder_.get(), fs(), observation_writer);
  }

  std::shared_ptr<const logger::ProjectContext> GetProjectContext() {
    return global_project_context_factory_->NewProjectContext(kCustomerId, kProjectId);
  }

 private:
  std::unique_ptr<logger::ProjectContextFactory> global_project_context_factory_;
  system_data::FakeSystemData system_data_;
  std::unique_ptr<MetadataBuilder> metadata_builder_;
};

TEST_F(LocalAggregationTest, AddEventWorks) {
  logger::testing::FakeObservationStore test_writer;
  logger::ObservationWriter observation_writer(&test_writer, nullptr);
  auto aggregation = MakeLocalAggregation(&observation_writer);

  ASSERT_THAT(fs()->ListFiles(absl::StrCat(local_aggregation_store_path(), "/", kCustomerId, "/",
                                           kProjectId))
                  .ConsumeValueOrDie(),
              Not(Contains(absl::StrCat(kOccurrenceMetricMetricId))));

  logger::EventRecord record(GetProjectContext(), kOccurrenceMetricMetricId);
  record.event()->mutable_occurrence_event();
  aggregation->AddEvent(record);

  ASSERT_THAT(fs()->ListFiles(absl::StrCat(local_aggregation_store_path(), "/", kCustomerId, "/",
                                           kProjectId))
                  .ConsumeValueOrDie(),
              Contains(absl::StrCat(kOccurrenceMetricMetricId)));
}

TEST_F(LocalAggregationTest, DisableWorks) {
  logger::testing::FakeObservationStore test_writer;
  logger::ObservationWriter observation_writer(&test_writer, nullptr);
  auto aggregation = MakeLocalAggregation(&observation_writer);
  aggregation->Disable(true);

  // Check that the no data is present.
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(local_aggregation_store_path(), "/", kCustomerId, "/",
                                           kProjectId))
                  .ConsumeValueOrDie(),
              Not(Contains(absl::StrCat(kOccurrenceMetricMetricId))));

  logger::EventRecord record(GetProjectContext(), kOccurrenceMetricMetricId);
  record.event()->mutable_occurrence_event();
  aggregation->AddEvent(record);

  // Check that the data still isn't present.
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(local_aggregation_store_path(), "/", kCustomerId, "/",
                                           kProjectId))
                  .ConsumeValueOrDie(),
              Not(Contains(absl::StrCat(kOccurrenceMetricMetricId))));
}

TEST_F(LocalAggregationTest, GenerateObservationsWorks) {
  auto project_context = GetProjectContext();
  auto metric = project_context->GetMetric(kOccurrenceMetricMetricId);
  // This test assumes that this metric has no reports with added privacy.
  for (const auto& report : metric->reports()) {
    ASSERT_EQ(ReportDefinition::NO_ADDED_PRIVACY, report.privacy_level());
  }

  logger::testing::FakeObservationStore test_writer;
  logger::ObservationWriter observation_writer(&test_writer, nullptr);
  auto aggregation = MakeLocalAggregation(&observation_writer);

  logger::EventRecord record(project_context, kOccurrenceMetricMetricId);
  record.event()->mutable_occurrence_event()->set_count(1);
  record.event()->set_day_index(1);
  record.event()->set_hour_id(1);
  ASSERT_TRUE(aggregation->AddEvent(record).ok());
  ASSERT_TRUE(aggregation->AddEvent(record).ok());
  ASSERT_TRUE(aggregation->AddEvent(record).ok());
  ASSERT_TRUE(aggregation->AddEvent(record).ok());

  aggregation->GenerateAggregatedObservations(TimeInfo::FromHourId(1), TimeInfo::FromHourId(1));

  ASSERT_EQ(test_writer.num_observations_added(), 2);
  ASSERT_EQ(test_writer.metadata_received[0]->customer_id(), kCustomerId);
  std::unique_ptr<observation_store::StoredObservation> last_stored_observation =
      std::move(test_writer.messages_received[0]);
  ASSERT_TRUE(last_stored_observation->has_unencrypted());

  Observation last_observation = last_stored_observation->unencrypted();
  ASSERT_TRUE(last_observation.has_integer());
  ASSERT_GT(last_observation.integer().values_size(), 0);
  ASSERT_EQ(last_observation.integer().values(0).value(), 4);
}

}  // namespace cobalt::local_aggregation