training_data_collector_impl_unittest.cc

// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/segmentation_platform/internal/data_collection/training_data_collector_impl.h"

#include <map>
#include <string_view>

#include "base/functional/callback_helpers.h"
#include "base/memory/raw_ptr.h"
#include "base/metrics/metrics_hashes.h"
#include "base/test/gmock_callback_support.h"
#include "base/test/metrics/histogram_tester.h"
#include "base/test/simple_test_clock.h"
#include "base/test/task_environment.h"
#include "base/time/time.h"
#include "components/prefs/testing_pref_service.h"
#include "components/segmentation_platform/internal/constants.h"
#include "components/segmentation_platform/internal/data_collection/training_data_collector.h"
#include "components/segmentation_platform/internal/database/cached_result_provider.h"
#include "components/segmentation_platform/internal/database/cached_result_writer.h"
#include "components/segmentation_platform/internal/database/client_result_prefs.h"
#include "components/segmentation_platform/internal/database/config_holder.h"
#include "components/segmentation_platform/internal/database/mock_signal_storage_config.h"
#include "components/segmentation_platform/internal/database/test_segment_info_database.h"
#include "components/segmentation_platform/internal/execution/processing/feature_list_query_processor.h"
#include "components/segmentation_platform/internal/execution/processing/mock_feature_list_query_processor.h"
#include "components/segmentation_platform/internal/metadata/metadata_utils.h"
#include "components/segmentation_platform/internal/mock_ukm_data_manager.h"
#include "components/segmentation_platform/internal/platform_options.h"
#include "components/segmentation_platform/internal/proto/model_prediction.pb.h"
#include "components/segmentation_platform/internal/segmentation_ukm_helper.h"
#include "components/segmentation_platform/internal/selection/segmentation_result_prefs.h"
#include "components/segmentation_platform/internal/signals/mock_histogram_signal_handler.h"
#include "components/segmentation_platform/internal/signals/mock_user_action_signal_handler.h"
#include "components/segmentation_platform/internal/stats.h"
#include "components/segmentation_platform/public/config.h"
#include "components/segmentation_platform/public/features.h"
#include "components/segmentation_platform/public/local_state_helper.h"
#include "components/segmentation_platform/public/model_provider.h"
#include "components/segmentation_platform/public/proto/model_metadata.pb.h"
#include "components/segmentation_platform/public/proto/segmentation_platform.pb.h"
#include "components/segmentation_platform/public/proto/types.pb.h"
#include "components/segmentation_platform/public/segmentation_platform_service.h"
#include "components/ukm/test_ukm_recorder.h"
#include "services/metrics/public/cpp/ukm_builders.h"
#include "services/metrics/public/cpp/ukm_source_id.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace segmentation_platform {
namespace {

RunOnceCallback;
RunOnceCallbackRepeatedly;
_;
NiceMock;
Return;
Segmentation_ModelExecution;

constexpr auto kTestOptimizationTarget0 = …;
constexpr auto kTestOptimizationTarget1 = …;
constexpr auto kTestOptimizationTarget2 = …;
constexpr char kHistogramName0[] = …;
constexpr char kHistogramName1[] = …;
constexpr char kSegmentationKey[] = …;
constexpr char kSegmentationKey2[] = …;
constexpr int64_t kModelVersion = …;
constexpr int kSample = …;
constexpr DecisionType kOnDemandDecisionType = …;
constexpr DecisionType kPeriodicDecisionType = …;

class MockModelManager : public ModelManager { … };

class TrainingDataCollectorImplTest
    : public ::testing::Test,
      public ::testing::WithParamInterface<bool> { … };

INSTANTIATE_TEST_SUITE_P(…);

// No segment info in database. Do nothing.
TEST_P(TrainingDataCollectorImplTest, NoSegment) { … }

// Histogram not in the output list will not trigger a training data report..
TEST_P(TrainingDataCollectorImplTest, IrrelevantHistogramNotReported) { … }

// No UKM report due to minimum data collection time not met.
TEST_P(TrainingDataCollectorImplTest, SignalCollectionRequirementNotMet) { … }

// No UKM report due to model updated recently.
TEST_P(TrainingDataCollectorImplTest, ModelUpdatedRecently) { … }

// No report if UKM is enabled recently.
TEST_P(TrainingDataCollectorImplTest, PartialOutputNotAllowed) { … }

// No training data recorded on startup if upload_tensor is set to false in
// continuous collection.
TEST_P(TrainingDataCollectorImplTest,
       ContinuousCollectionOnStartupWithoutUploadTensor) { … }

// Tests that continuous collection happens on startup.
TEST_P(TrainingDataCollectorImplTest, ContinuousCollectionOnStartupNoDelay) { … }

// Tests that continuous collection do not collect for ondemand models on
// startup.
TEST_P(TrainingDataCollectorImplTest,
       OnDemandModelsDoNotTriggerPeriodicCollection) { … }

// Tests that ReportCollectedContinuousTrainingData() works well later if
// no data is reported on start up.
TEST_P(TrainingDataCollectorImplTest,
       ReportCollectedContinuousTrainingData_LegacyConfig) { … }

// Tests that ReportCollectedContinuousTrainingData() works well later if
// no data is reported on start up.
TEST_P(TrainingDataCollectorImplTest,
       ReportCollectedContinuousTrainingData_MultiOutputConfig) { … }

TEST_P(TrainingDataCollectorImplTest, ContinuousWithExactPredictionNotSet) { … }

TEST_P(TrainingDataCollectorImplTest, ContinuousWithExactPrediction) { … }

TEST_P(TrainingDataCollectorImplTest, ContinuousWithFlexibleObservation) { … }

TEST_P(TrainingDataCollectorImplTest, ContinuousWithDelay) { … }

// Tests that after a data collection, another data collection won't happen
// immediately afterwards.
TEST_P(TrainingDataCollectorImplTest,
       NoImmediateDataCollectionAfterLastCollection) { … }

// Tests that if UKM allowed timestamp is not set in local state, data
// collection won't happen.
TEST_P(TrainingDataCollectorImplTest, NoDataCollectionIfUkmAllowedPrefNotSet) { … }

// Tests that if uma histogram trigger is set, collection will happen when the
// trigger histogram is observed.
TEST_P(TrainingDataCollectorImplTest, DataCollectionWithEnumHistogramTrigger) { … }

// Tests that if uma user action trigger is set, collection will happen when the
// trigger user action is observed.
TEST_P(TrainingDataCollectorImplTest, DataCollectionWithUserActionTrigger) { … }

// Tests that if uma user action trigger is set, collection will not happen
// without upload_tensor = true.
TEST_P(TrainingDataCollectorImplTest,
       DataCollectionWithTriggerWithoutUploadTensor) { … }

// A histogram interested by multiple model will trigger multiple UKM reports.
TEST_P(TrainingDataCollectorImplTest,
       DataCollectionWithUMATrigger_MultipleModels) { … }

// Tests that if no uma histogram trigger is set, collection will happen when
// the time delay passes.
TEST_P(TrainingDataCollectorImplTest, DataCollectionWithTimeTrigger) { … }

TEST_P(TrainingDataCollectorImplTest, DataCollectionWithStoreToDisk) { … }

TEST_P(TrainingDataCollectorImplTest, DataCollectionWithTriggerAPI) { … }

// No training data recorded if upload_tensor is set to false in on-demand
// collection using trigger API.
TEST_P(TrainingDataCollectorImplTest,
       DataCollectionTriggerAPIWithoutUploadTensor) { … }

TEST_P(TrainingDataCollectorImplTest,
       DataCollectionWithTriggerAPIForPreferredSegment) { … }

// Tests that we don't collect training data if input processing fails.
TEST_P(TrainingDataCollectorImplTest,
       DataCollectionSkippedWhenInputProcessingFails) { … }

}  // namespace
}  // namespace segmentation_platform
chromium/components/segmentation_platform/internal/data_collection/training_data_collector_impl_unittest.cc
