page_content_annotation_job_executor_unittest.cc 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. // Copyright 2021 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "components/optimization_guide/core/page_content_annotation_job_executor.h"
  5. #include "base/test/scoped_feature_list.h"
  6. #include "base/test/scoped_run_loop_timeout.h"
  7. #include "base/test/task_environment.h"
  8. #include "components/optimization_guide/core/optimization_guide_features.h"
  9. #include "components/optimization_guide/proto/models.pb.h"
  10. #include "testing/gmock/include/gmock/gmock.h"
  11. #include "testing/gtest/include/gtest/gtest.h"
  12. namespace optimization_guide {
  13. namespace {
  14. const std::vector<WeightedIdentifier> kOutput{WeightedIdentifier(1337, 1.0)};
  15. }
  16. class TestJobExecutor : public PageContentAnnotationJobExecutor {
  17. public:
  18. TestJobExecutor() = default;
  19. virtual ~TestJobExecutor() = default;
  20. // PageContentAnnotationJobExecutor:
  21. void ExecuteOnSingleInput(
  22. AnnotationType annotation_type,
  23. const std::string& input,
  24. base::OnceCallback<void(const BatchAnnotationResult&)> callback)
  25. override {
  26. std::move(callback).Run(
  27. BatchAnnotationResult::CreatePageTopicsResult(input, kOutput));
  28. }
  29. };
  30. class PageContentAnnotationJobExecutorTest : public testing::Test {
  31. public:
  32. PageContentAnnotationJobExecutorTest() {
  33. scoped_feature_list_.InitAndEnableFeature(
  34. features::kPageContentAnnotations);
  35. }
  36. ~PageContentAnnotationJobExecutorTest() override = default;
  37. private:
  38. base::test::TaskEnvironment task_environment_;
  39. base::test::ScopedFeatureList scoped_feature_list_;
  40. };
  41. TEST_F(PageContentAnnotationJobExecutorTest, FullFlow) {
  42. TestJobExecutor job_executor;
  43. std::vector<BatchAnnotationResult> results;
  44. // This callback will be run before the run loop quit closure below.
  45. BatchAnnotationCallback outside_callers_result_callback = base::BindOnce(
  46. [](std::vector<BatchAnnotationResult>* out_results,
  47. const std::vector<BatchAnnotationResult>& in_results) {
  48. *out_results = in_results;
  49. },
  50. &results);
  51. std::unique_ptr<PageContentAnnotationJob> job =
  52. std::make_unique<PageContentAnnotationJob>(
  53. std::move(outside_callers_result_callback),
  54. std::vector<std::string>{"input1", "input2"},
  55. AnnotationType::kPageTopics);
  56. // Actual model execution can take a little while, so try to keep tests from
  57. // flaking.
  58. base::test::ScopedRunLoopTimeout scoped_timeout(FROM_HERE, base::Seconds(60));
  59. base::RunLoop run_loop;
  60. job_executor.ExecuteJob(run_loop.QuitClosure(), std::move(job));
  61. run_loop.Run();
  62. ASSERT_EQ(2U, results.size());
  63. EXPECT_EQ(results[0].input(), "input1");
  64. EXPECT_EQ(results[0].topics(), absl::make_optional(kOutput));
  65. EXPECT_EQ(results[1].input(), "input2");
  66. EXPECT_EQ(results[1].topics(), absl::make_optional(kOutput));
  67. }
  68. } // namespace optimization_guide