| // Copyright 2019 The Chromium 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 <memory> |
| #include <utility> |
| #include <vector> |
| |
| #include "base/bind.h" |
| #include "base/test/task_environment.h" |
| #include "base/threading/sequenced_task_runner_handle.h" |
| #include "media/learning/impl/learning_task_controller_helper.h" |
| #include "testing/gtest/include/gtest/gtest.h" |
| |
| namespace media { |
| namespace learning { |
| |
| class LearningTaskControllerHelperTest : public testing::Test { |
| public: |
| class FakeFeatureProvider : public FeatureProvider { |
| public: |
| FakeFeatureProvider(FeatureVector* features_out, |
| FeatureProvider::FeatureVectorCB* cb_out) |
| : features_out_(features_out), cb_out_(cb_out) {} |
| |
| // Do nothing, except note that we were called. |
| void AddFeatures(FeatureVector features, |
| FeatureProvider::FeatureVectorCB cb) override { |
| *features_out_ = std::move(features); |
| *cb_out_ = std::move(cb); |
| } |
| |
| FeatureVector* features_out_; |
| FeatureProvider::FeatureVectorCB* cb_out_; |
| }; |
| |
| LearningTaskControllerHelperTest() { |
| task_runner_ = base::SequencedTaskRunnerHandle::Get(); |
| task_.name = "example_task"; |
| |
| example_.features.push_back(FeatureValue(1)); |
| example_.features.push_back(FeatureValue(2)); |
| example_.features.push_back(FeatureValue(3)); |
| example_.target_value = TargetValue(123); |
| example_.weight = 100u; |
| |
| id_ = base::UnguessableToken::Create(); |
| } |
| |
| ~LearningTaskControllerHelperTest() override { |
| // To prevent a memory leak, reset the helper. This will post destruction |
| // of other objects, so RunUntilIdle(). |
| helper_.reset(); |
| task_environment_.RunUntilIdle(); |
| } |
| |
| void CreateClient(bool include_fp) { |
| // Create the fake feature provider, and get a pointer to it. |
| base::SequenceBound<FakeFeatureProvider> sb_fp; |
| if (include_fp) { |
| sb_fp = base::SequenceBound<FakeFeatureProvider>(task_runner_, |
| &fp_features_, &fp_cb_); |
| task_environment_.RunUntilIdle(); |
| } |
| |
| // TODO(liberato): make sure this works without a fp. |
| helper_ = std::make_unique<LearningTaskControllerHelper>( |
| task_, |
| base::BindRepeating( |
| &LearningTaskControllerHelperTest::OnLabelledExample, |
| base::Unretained(this)), |
| std::move(sb_fp)); |
| } |
| |
| void OnLabelledExample(LabelledExample example, ukm::SourceId source_id) { |
| most_recent_example_ = std::move(example); |
| most_recent_source_id_ = source_id; |
| } |
| |
| // Since we're friends but the tests aren't. |
| size_t pending_example_count() const { |
| return helper_->pending_example_count_for_testing(); |
| } |
| |
| base::test::TaskEnvironment task_environment_; |
| |
| scoped_refptr<base::SequencedTaskRunner> task_runner_; |
| |
| std::unique_ptr<LearningTaskControllerHelper> helper_; |
| |
| // Most recent features / cb given to our FakeFeatureProvider. |
| FeatureVector fp_features_; |
| FeatureProvider::FeatureVectorCB fp_cb_; |
| |
| // Most recently added example via OnLabelledExample, if any. |
| absl::optional<LabelledExample> most_recent_example_; |
| ukm::SourceId most_recent_source_id_; |
| |
| LearningTask task_; |
| |
| base::UnguessableToken id_; |
| |
| LabelledExample example_; |
| }; |
| |
| TEST_F(LearningTaskControllerHelperTest, AddingAnExampleWithoutFPWorks) { |
| // A helper that doesn't use a FeatureProvider should forward examples as soon |
| // as they're done. |
| CreateClient(false); |
| ukm::SourceId source_id = 2; |
| helper_->BeginObservation(id_, example_.features, source_id); |
| EXPECT_EQ(pending_example_count(), 1u); |
| helper_->CompleteObservation( |
| id_, ObservationCompletion(example_.target_value, example_.weight)); |
| EXPECT_TRUE(most_recent_example_); |
| EXPECT_EQ(*most_recent_example_, example_); |
| EXPECT_EQ(most_recent_example_->weight, example_.weight); |
| EXPECT_EQ(most_recent_source_id_, source_id); |
| EXPECT_EQ(pending_example_count(), 0u); |
| } |
| |
| TEST_F(LearningTaskControllerHelperTest, DropTargetValueWithoutFPWorks) { |
| // Verify that we can drop an example without labelling it. |
| CreateClient(false); |
| helper_->BeginObservation(id_, example_.features, absl::nullopt); |
| EXPECT_EQ(pending_example_count(), 1u); |
| helper_->CancelObservation(id_); |
| task_environment_.RunUntilIdle(); |
| EXPECT_FALSE(most_recent_example_); |
| EXPECT_EQ(pending_example_count(), 0u); |
| } |
| |
| TEST_F(LearningTaskControllerHelperTest, AddTargetValueBeforeFP) { |
| // Verify that an example is added if the target value arrives first. |
| CreateClient(true); |
| helper_->BeginObservation(id_, example_.features, absl::nullopt); |
| EXPECT_EQ(pending_example_count(), 1u); |
| task_environment_.RunUntilIdle(); |
| // The feature provider should know about the example. |
| EXPECT_EQ(fp_features_, example_.features); |
| |
| // Add the targe value and verify that the example wasn't added yet. |
| helper_->CompleteObservation( |
| id_, ObservationCompletion(example_.target_value, example_.weight)); |
| EXPECT_FALSE(most_recent_example_); |
| EXPECT_EQ(pending_example_count(), 1u); |
| |
| // Add the features, and verify that they arrive at the AddExampleCB. |
| example_.features[0] = FeatureValue(456); |
| std::move(fp_cb_).Run(example_.features); |
| task_environment_.RunUntilIdle(); |
| EXPECT_EQ(pending_example_count(), 0u); |
| EXPECT_TRUE(most_recent_example_); |
| EXPECT_EQ(*most_recent_example_, example_); |
| EXPECT_EQ(most_recent_example_->weight, example_.weight); |
| } |
| |
| TEST_F(LearningTaskControllerHelperTest, DropTargetValueBeforeFP) { |
| // Verify that an example is correctly dropped before the FP adds features. |
| CreateClient(true); |
| helper_->BeginObservation(id_, example_.features, absl::nullopt); |
| EXPECT_EQ(pending_example_count(), 1u); |
| task_environment_.RunUntilIdle(); |
| // The feature provider should know about the example. |
| EXPECT_EQ(fp_features_, example_.features); |
| |
| // Cancel the observation. |
| helper_->CancelObservation(id_); |
| // We don't care if the example is still queued or not, only that we can |
| // add features and have it be zero by then. |
| |
| // Add the features, and verify that the pending example is removed and no |
| // example was sent to us. |
| example_.features[0] = FeatureValue(456); |
| std::move(fp_cb_).Run(example_.features); |
| task_environment_.RunUntilIdle(); |
| EXPECT_EQ(pending_example_count(), 0u); |
| EXPECT_FALSE(most_recent_example_); |
| } |
| |
| TEST_F(LearningTaskControllerHelperTest, AddTargetValueAfterFP) { |
| // Verify that an example is added if the target value arrives second. |
| CreateClient(true); |
| helper_->BeginObservation(id_, example_.features, absl::nullopt); |
| EXPECT_EQ(pending_example_count(), 1u); |
| task_environment_.RunUntilIdle(); |
| // The feature provider should know about the example. |
| EXPECT_EQ(fp_features_, example_.features); |
| EXPECT_EQ(pending_example_count(), 1u); |
| |
| // Add the features, and verify that the example isn't sent yet. |
| example_.features[0] = FeatureValue(456); |
| std::move(fp_cb_).Run(example_.features); |
| task_environment_.RunUntilIdle(); |
| EXPECT_FALSE(most_recent_example_); |
| EXPECT_EQ(pending_example_count(), 1u); |
| |
| // Add the targe value and verify that the example is added. |
| helper_->CompleteObservation( |
| id_, ObservationCompletion(example_.target_value, example_.weight)); |
| EXPECT_TRUE(most_recent_example_); |
| EXPECT_EQ(*most_recent_example_, example_); |
| EXPECT_EQ(most_recent_example_->weight, example_.weight); |
| EXPECT_EQ(pending_example_count(), 0u); |
| } |
| |
| TEST_F(LearningTaskControllerHelperTest, DropTargetValueAfterFP) { |
| // Verify that we can cancel the observationc after sending features. |
| CreateClient(true); |
| helper_->BeginObservation(id_, example_.features, absl::nullopt); |
| EXPECT_EQ(pending_example_count(), 1u); |
| task_environment_.RunUntilIdle(); |
| // The feature provider should know about the example. |
| EXPECT_EQ(fp_features_, example_.features); |
| EXPECT_EQ(pending_example_count(), 1u); |
| |
| // Add the features, and verify that the example isn't sent yet. We do care |
| // that the example is still pending, since we haven't actually dropped the |
| // callback yet; we might send a TargetValue. |
| example_.features[0] = FeatureValue(456); |
| std::move(fp_cb_).Run(example_.features); |
| task_environment_.RunUntilIdle(); |
| EXPECT_FALSE(most_recent_example_); |
| EXPECT_EQ(pending_example_count(), 1u); |
| |
| // Cancel the observation, and verify that the pending example has been |
| // removed, and no example was sent to us. |
| helper_->CancelObservation(id_); |
| task_environment_.RunUntilIdle(); |
| EXPECT_FALSE(most_recent_example_); |
| EXPECT_EQ(pending_example_count(), 0u); |
| } |
| |
| } // namespace learning |
| } // namespace media |