src/local_aggregation_1_1/aggregation_procedures/testing/test_aggregation_procedure.h

// 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.

#ifndef COBALT_SRC_LOCAL_AGGREGATION_1_1_AGGREGATION_PROCEDURES_TESTING_TEST_AGGREGATION_PROCEDURE_H_
#define COBALT_SRC_LOCAL_AGGREGATION_1_1_AGGREGATION_PROCEDURES_TESTING_TEST_AGGREGATION_PROCEDURE_H_

#include <gtest/gtest.h>

#include "src/lib/util/datetime_util.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/project_context_factory.h"
#include "src/public/lib/registry_identifiers.h"

namespace cobalt::local_aggregation::testing {

class TestAggregationProcedure : public util::testing::TestWithFiles {
 public:
  TestAggregationProcedure() {
    global_project_context_factory_ =
        logger::ProjectContextFactory::CreateFromCobaltRegistryBase64(kCobaltRegistryBase64);
  }

  std::unique_ptr<logger::ProjectContext> GetProjectContext() {
    return global_project_context_factory_->NewProjectContext(
        lib::CustomerIdentifier(kCustomerId).ForProject(kProjectId));
  }

  MetricDefinition GetMetricDef(uint32_t metric_id) {
    current_metric_id_ = metric_id;

    auto project_context = global_project_context_factory_->NewProjectContext(
        lib::CustomerIdentifier(kCustomerId).ForProject(kProjectId));
    return *project_context->GetMetric(metric_id);
  }

  ReportDefinition GetReportDef(uint32_t metric_id, int report_index) {
    current_metric_id_ = metric_id;
    current_report_index_ = report_index;

    return GetMetricDef(metric_id).reports(report_index);
  }

  std::unique_ptr<AggregationProcedure> GetProcedureFor(uint32_t metric_id, int report_index) {
    current_metric_id_ = metric_id;
    current_report_index_ = report_index;

    const MetricDefinition& metric = GetMetricDef(current_metric_id_);
    const ReportDefinition& report = metric.reports(current_report_index_);
    current_event_vector_buffer_max_ = report.event_vector_buffer_max();
    current_string_buffer_max_ = report.string_buffer_max();

    auto agg_or = AggregationProcedure::Get(metric, report);
    EXPECT_TRUE(agg_or.ok());
    return std::move(agg_or).value();
  }

  std::unique_ptr<logger::EventRecord> MakeEventRecord(util::TimeInfo time_info) {
    auto record = logger::EventRecord::MakeEventRecord(
                      global_project_context_factory_->NewProjectContext(
                          lib::CustomerIdentifier(kCustomerId).ForProject(kProjectId)),
                      current_metric_id_)
                      .value();

    record->event()->set_day_index(time_info.day_index);
    record->event()->set_hour_id(time_info.hour_id);

    return record;
  }

  void AddOccurrenceEventsForDay(
      uint32_t num_events, uint32_t day_index, uint64_t system_profile_hash,
      AggregationProcedure& procedure, ReportAggregate& aggregate,
      std::optional<std::chrono::system_clock::time_point> event_time = std::nullopt) {
    AddOccurrenceEventsForDayWithCount(num_events, 1, day_index, system_profile_hash, procedure,
                                       aggregate, event_time);
  }

  void AddOccurrenceEventsForDayWithCount(
      uint32_t num_events, uint64_t count, uint32_t day_index, uint64_t system_profile_hash,
      AggregationProcedure& procedure, ReportAggregate& aggregate,
      std::optional<std::chrono::system_clock::time_point> event_time = std::nullopt) {
    auto record = MakeEventRecord(util::TimeInfo::FromDayIndex(day_index));
    OccurrenceEvent* event = record->event()->mutable_occurrence_event();
    event->add_event_code(0);
    event->set_count(count);

    for (uint32_t i = 1; i <= num_events; i++) {
      event->set_event_code(0, i);
      for (uint32_t j = 0; j < i; j++) {
        procedure.UpdateAggregate(
            *record, aggregate, system_profile_hash,
            event_time.value_or(util::FromUnixSeconds(util::DayIndexToUnixSeconds(day_index))));
      }
    }
  }

  void AddOccurrenceEventsForHour(
      uint32_t num_events, uint32_t hour_id, uint64_t system_profile_hash,
      AggregationProcedure& procedure, ReportAggregate& aggregate,
      std::optional<std::chrono::system_clock::time_point> event_time = std::nullopt) {
    AddOccurrenceEventsForHourWithCount(num_events, 1, hour_id, system_profile_hash, procedure,
                                        aggregate, event_time);
  }

  void AddOccurrenceEventsForHourWithCount(
      uint32_t num_events, uint64_t count, uint32_t hour_id, uint64_t system_profile_hash,
      AggregationProcedure& procedure, ReportAggregate& aggregate,
      std::optional<std::chrono::system_clock::time_point> event_time = std::nullopt) {
    auto record = MakeEventRecord(util::TimeInfo::FromHourId(hour_id));
    OccurrenceEvent* event = record->event()->mutable_occurrence_event();
    event->add_event_code(0);
    event->set_count(count);

    for (uint32_t i = 1; i <= num_events; i++) {
      event->set_event_code(0, i);
      for (uint32_t j = 0; j < i; j++) {
        procedure.UpdateAggregate(
            *record, aggregate, system_profile_hash,
            event_time.value_or(util::FromUnixSeconds(util::HourIdToUnixSeconds(hour_id))));
      }
    }
  }

  void AddIntegerEvents(
      uint32_t num_events, int64_t value, uint32_t hour_id, uint64_t system_profile_hash,
      AggregationProcedure& procedure, ReportAggregate& aggregate,
      std::optional<std::chrono::system_clock::time_point> event_time = std::nullopt) {
    auto record = MakeEventRecord(util::TimeInfo::FromHourId(hour_id));
    IntegerEvent* event = record->event()->mutable_integer_event();
    event->add_event_code(0);
    event->set_value(value);

    for (uint32_t i = 1; i <= num_events; i++) {
      event->set_event_code(0, i);
      for (uint32_t j = 0; j < i; j++) {
        procedure.UpdateAggregate(
            *record, aggregate, system_profile_hash,
            event_time.value_or(util::FromUnixSeconds(util::HourIdToUnixSeconds(hour_id))));
      }
    }
  }

  [[nodiscard]] uint64_t GetEventVectorBufferMax() const {
    return current_event_vector_buffer_max_;
  }

  [[nodiscard]] uint32_t GetStringBufferMax() const { return current_string_buffer_max_; }

 private:
  std::unique_ptr<logger::ProjectContextFactory> global_project_context_factory_;
  uint32_t current_metric_id_;
  int current_report_index_;
  uint64_t current_event_vector_buffer_max_;
  uint32_t current_string_buffer_max_;
};

}  // namespace cobalt::local_aggregation::testing

#endif  // COBALT_SRC_LOCAL_AGGREGATION_1_1_AGGREGATION_PROCEDURES_TESTING_TEST_AGGREGATION_PROCEDURE_H_