ReportQueueContext

Bug: b:159750481
Change-Id: Ie159965cb3911fd6c4effeb601f98af7ab517114
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2616321
Commit-Queue: Zach Trudo <zatrudo@google.com>
Reviewed-by: Leonid Baraz <lbaraz@chromium.org>
Cr-Commit-Position: refs/heads/master@{#845926}

chrome/browser/BUILD.gn

@@ -1242,6 +1242,8 @@ static_library("browser") {
     "policy/messaging_layer/upload/upload_client.h",
     "policy/messaging_layer/util/backoff_settings.cc",
     "policy/messaging_layer/util/backoff_settings.h",
+    "policy/messaging_layer/util/report_queue_manual_test_context.cc",
+    "policy/messaging_layer/util/report_queue_manual_test_context.h",
     "policy/messaging_layer/util/shared_queue.h",
     "policy/messaging_layer/util/shared_vector.h",
     "policy/messaging_layer/util/status.cc",

chrome/browser/policy/messaging_layer/util/report_queue_manual_test_context.cc

+// Copyright 2021 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 "chrome/browser/policy/messaging_layer/util/report_queue_manual_test_context.h"
+
+#include "base/callback.h"
+#include "base/no_destructor.h"
+#include "base/optional.h"
+#include "base/sequenced_task_runner.h"
+#include "base/strings/string_number_conversions.h"
+#include "base/time/time.h"
+#include "chrome/browser/policy/dm_token_utils.h"
+#include "chrome/browser/policy/messaging_layer/public/report_client.h"
+#include "chrome/browser/policy/messaging_layer/public/report_queue.h"
+#include "chrome/browser/policy/messaging_layer/public/report_queue_configuration.h"
+#include "chrome/browser/policy/messaging_layer/util/status.h"
+#include "chrome/browser/policy/messaging_layer/util/task_runner_context.h"
+#include "chrome/browser/profiles/profile_manager.h"
+#include "components/policy/core/common/cloud/dm_token.h"
+#include "components/policy/proto/record_constants.pb.h"
+#include "content/public/browser/browser_task_traits.h"
+#include "content/public/browser/browser_thread.h"
+
+namespace reporting {
+
+ReportQueueManualTestContext::ReportQueueManualTestContext(
+    base::TimeDelta frequency,
+    uint64_t number_of_messages_to_enqueue,
+    Destination destination,
+    Priority priority,
+    CompletionCallback completion_cb,
+    scoped_refptr<base::SequencedTaskRunner> sequenced_task_runner)
+    : TaskRunnerContext<Status>(std::move(completion_cb),
+                                sequenced_task_runner),
+      frequency_(frequency),
+      number_of_messages_to_enqueue_(number_of_messages_to_enqueue),
+      destination_(destination),
+      priority_(priority) {}
+
+ReportQueueManualTestContext::~ReportQueueManualTestContext() = default;
+
+void ReportQueueManualTestContext::SetBuildReportQueueCallbackForTests(
+    BuildReportQueueCallback build_report_queue_cb) {
+  *GetBuildReportQueueCallback() = std::move(build_report_queue_cb);
+}
+
+void ReportQueueManualTestContext::OnStart() {
+  if (destination_ == UNDEFINED_DESTINATION) {
+    Status invalid_destination = Status(
+        error::INVALID_ARGUMENT, "Destination was UNDEFINED_DESTINATION");
+    LOG(ERROR) << invalid_destination;
+    Complete(invalid_destination);
+    return;
+  }
+
+  if (priority_ == UNDEFINED_PRIORITY) {
+    Status invalid_priority = Status(error::INVALID_ARGUMENT,
+                                     "Destination was UNDEFINED_DESTINATION");
+    LOG(ERROR) << invalid_priority;
+    Complete(invalid_priority);
+    return;
+  }
+
+  // The DMToken must be retrieved on the UI thread.
+  base::PostTask(FROM_HERE, {content::BrowserThread::UI},
+                 base::BindOnce(&ReportQueueManualTestContext::GetDmToken,
+                                base::Unretained(this)));
+}
+
+void ReportQueueManualTestContext::GetDmToken() {
+  const policy::DMToken dm_token =
+      policy::GetDMToken(ProfileManager::GetPrimaryUserProfile());
+  Schedule(&ReportQueueManualTestContext::OnDmTokenResponse,
+           base::Unretained(this), dm_token);
+}
+
+void ReportQueueManualTestContext::OnDmTokenResponse(policy::DMToken dm_token) {
+  CheckOnValidSequence();
+  if (!dm_token.is_valid()) {
+    Status invalid_dm_token =
+        Status(error::FAILED_PRECONDITION, "Cannot retrieve a valid DMToken");
+    LOG(ERROR) << invalid_dm_token;
+    Complete(invalid_dm_token);
+    return;
+  }
+  dm_token_ = dm_token;
+  BuildReportQueue();
+}
+
+void ReportQueueManualTestContext::BuildReportQueue() {
+  CheckOnValidSequence();
+  ReportQueueConfiguration::PolicyCheckCallback policy_check_cb =
+      base::BindRepeating([]() -> Status { return Status::StatusOK(); });
+  auto config_result = reporting::ReportQueueConfiguration::Create(
+      dm_token_, destination_, std::move(policy_check_cb));
+  if (!config_result.ok()) {
+    Complete(config_result.status());
+    return;
+  }
+
+  if (*GetBuildReportQueueCallback()) {
+    std::move(*GetBuildReportQueueCallback())
+        .Run(
+            std::move(config_result.ValueOrDie()),
+            base::BindOnce(&ReportQueueManualTestContext::OnReportQueueResponse,
+                           base::Unretained(this)));
+    return;
+  }
+
+  reporting::ReportingClient::CreateReportQueue(
+      std::move(config_result.ValueOrDie()),
+      base::BindOnce(&ReportQueueManualTestContext::OnReportQueueResponse,
+                     base::Unretained(this)));
+}
+
+void ReportQueueManualTestContext::OnReportQueueResponse(
+    StatusOr<std::unique_ptr<ReportQueue>> report_queue_result) {
+  if (!report_queue_result.ok()) {
+    Complete(report_queue_result.status());
+    return;
+  }
+
+  report_queue_ = std::move(report_queue_result.ValueOrDie());
+  Schedule(&ReportQueueManualTestContext::ScheduleEnqueue,
+           base::Unretained(this));
+}
+
+void ReportQueueManualTestContext::ScheduleEnqueue() {
+  CheckOnValidSequence();
+  if (number_of_messages_to_enqueue_ > 0u &&
+      number_of_messages_to_enqueue_ == number_of_enqueued_messages_) {
+    Complete(Status::StatusOK());
+    return;
+  }
+
+  ScheduleAfter(frequency_, &ReportQueueManualTestContext::Enqueue,
+                base::Unretained(this));
+}
+
+void ReportQueueManualTestContext::Enqueue() {
+  CheckOnValidSequence();
+  report_queue_->Enqueue(
+      base::NumberToString(value_++), priority_,
+      base::BindOnce(&ReportQueueManualTestContext::OnEnqueue,
+                     base::Unretained(this)));
+  number_of_enqueued_messages_++;
+}
+
+void ReportQueueManualTestContext::OnEnqueue(Status status) {
+  if (!status.ok()) {
+    Complete(status);
+    return;
+  }
+
+  Schedule(&ReportQueueManualTestContext::ScheduleEnqueue,
+           base::Unretained(this));
+}
+
+void ReportQueueManualTestContext::Complete(Status status) {
+  Schedule(&ReportQueueManualTestContext::Response, base::Unretained(this),
+           status);
+}
+
+// static
+ReportQueueManualTestContext::BuildReportQueueCallback*
+ReportQueueManualTestContext::GetBuildReportQueueCallback() {
+  static base::NoDestructor<BuildReportQueueCallback> callback{
+      BuildReportQueueCallback()};
+  return callback.get();
+}
+
+}  // namespace reporting

chrome/browser/policy/messaging_layer/util/report_queue_manual_test_context.h

+// Copyright 2021 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.
+
+#ifndef CHROME_BROWSER_POLICY_MESSAGING_LAYER_UTIL_REPORT_QUEUE_MANUAL_TEST_CONTEXT_H_
+#define CHROME_BROWSER_POLICY_MESSAGING_LAYER_UTIL_REPORT_QUEUE_MANUAL_TEST_CONTEXT_H_
+
+#include "base/callback.h"
+#include "base/sequenced_task_runner.h"
+#include "base/time/time.h"
+#include "chrome/browser/policy/messaging_layer/public/report_queue.h"
+#include "chrome/browser/policy/messaging_layer/util/status.h"
+#include "chrome/browser/policy/messaging_layer/util/task_runner_context.h"
+#include "components/policy/core/common/cloud/dm_token.h"
+#include "components/policy/proto/record_constants.pb.h"
+
+namespace reporting {
+
+// This is a test fixture for manually testing uploading of events. Should only
+// be used when debugging or testing. It will enqueue an incrementing uint64
+// every |frequency| for the provided destination. All messages will be sent as
+// FAST_BATCH.
+//
+// This context can be used in the following way:
+//   Start<ReportQueueManualTestContext>(base::TimeDelta::FromSeconds(1),
+//                             /*number_of_messages_to_enqueue=*/10,
+//                             UPLOAD_EVENTS,
+//                             FAST_BATCH,
+//                             base::BindOnce(
+//                               [](Status status) {
+//                                 LOG(INFO) << status;
+//                              }),
+//                              base::ThreadPool::CreateSequencedTaskRunner(
+//                                base::TaskTraits()));
+// As configured this context will create a ReportQueue and upload an event to
+// the FAST_BATCH priority every second for 10 seconds.
+class ReportQueueManualTestContext : public TaskRunnerContext<Status> {
+ public:
+  using CompletionCallback = base::OnceCallback<void(Status)>;
+  using BuildReportQueueCallback = base::OnceCallback<void(
+      std::unique_ptr<ReportQueueConfiguration>,
+      base::OnceCallback<void(StatusOr<std::unique_ptr<ReportQueue>>)>)>;
+
+  ReportQueueManualTestContext(
+      base::TimeDelta frequency,
+      uint64_t number_of_messages_to_enqueue,
+      Destination destination,
+      Priority priorty,
+      CompletionCallback completion_cb,
+      scoped_refptr<base::SequencedTaskRunner> sequenced_task_runner);
+
+  // Sets the callback for building the report queue. Should be called before
+  // calling ReportQueueManualTestContext().
+  static void SetBuildReportQueueCallbackForTests(
+      BuildReportQueueCallback build_report_queue_cb);
+
+ private:
+  ~ReportQueueManualTestContext() override;
+
+  void OnStart() override;
+
+  void GetDmToken();
+  void OnDmTokenResponse(policy::DMToken dm_token);
+  void BuildReportQueue();
+  void OnReportQueueResponse(
+      StatusOr<std::unique_ptr<ReportQueue>> report_queue_result);
+  void ScheduleEnqueue();
+  void Enqueue();
+  void OnEnqueue(Status status);
+
+  void Complete(Status status);
+
+  static BuildReportQueueCallback* GetBuildReportQueueCallback();
+
+  // Frequency at which messages should be enqueued.
+  const base::TimeDelta frequency_;
+
+  // Total number of messages that should be enqueued.
+  // If set to 0, messages will be sent until the user ends the session or
+  // ReportQueue::Enqueue sends back a not-OK error.
+  const uint64_t number_of_messages_to_enqueue_;
+
+  const Destination destination_;
+  const Priority priority_;
+
+  // Counter for the number of messages sent. Should only be accessed while on
+  // sequence.
+  uint64_t number_of_enqueued_messages_{0u};
+
+  // Counter for the current value being enqueued. Should only be accessed while
+  // on sequence.
+  uint64_t value_{0u};
+
+  policy::DMToken dm_token_;
+  std::unique_ptr<ReportQueue> report_queue_;
+};
+
+}  // namespace reporting
+
+#endif  // CHROME_BROWSER_POLICY_MESSAGING_LAYER_UTIL_REPORT_QUEUE_MANUAL_TEST_CONTEXT_H_

chrome/browser/policy/messaging_layer/util/report_queue_manual_test_context_unittest.cc

+// Copyright 2021 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 "chrome/browser/policy/messaging_layer/util/report_queue_manual_test_context.h"
+
+#include "base/task_runner.h"
+#include "base/test/task_environment.h"
+#include "chrome/browser/policy/dm_token_utils.h"
+#include "chrome/browser/policy/messaging_layer/public/mock_report_queue.h"
+#include "chrome/browser/policy/messaging_layer/public/report_queue.h"
+#include "chrome/browser/policy/messaging_layer/util/status.h"
+#include "components/policy/core/common/cloud/dm_token.h"
+#include "components/policy/proto/record_constants.pb.h"
+#include "content/public/test/browser_task_environment.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace reporting {
+namespace {
+
+using testing::_;
+using testing::Invoke;
+using testing::WithArgs;
+
+class TestCallbackWaiter {
+ public:
+  TestCallbackWaiter() = default;
+
+  virtual void Signal() { run_loop_.Quit(); }
+
+  void Wait() { run_loop_.Run(); }
+
+ protected:
+  base::RunLoop run_loop_;
+};
+
+class TestCallbackWaiterWithCounter : public TestCallbackWaiter {
+ public:
+  explicit TestCallbackWaiterWithCounter(size_t counter_limit)
+      : counter_limit_(counter_limit) {}
+
+  void Signal() override {
+    DCHECK_GT(counter_limit_, 0u);
+    if (--counter_limit_ == 0u) {
+      run_loop_.Quit();
+    }
+  }
+
+ private:
+  std::atomic<size_t> counter_limit_;
+};
+
+class ReportQueueManualTestContextTest : public testing::Test {
+ protected:
+  ReportQueueManualTestContextTest() = default;
+
+  void SetUp() override {
+    auto mock_report_queue = std::make_unique<reporting::MockReportQueue>();
+    mock_report_queue_ = mock_report_queue.get();
+
+    auto build_report_queue_cb = base::BindOnce(
+        [](std::unique_ptr<ReportQueue> report_queue,
+           std::unique_ptr<ReportQueueConfiguration> report_queue_config,
+           base::OnceCallback<void(StatusOr<std::unique_ptr<ReportQueue>>)>
+               report_queue_cb) {
+          std::move(report_queue_cb).Run(std::move(report_queue));
+        },
+        std::move(mock_report_queue));
+
+    ReportQueueManualTestContext::SetBuildReportQueueCallbackForTests(
+        std::move(build_report_queue_cb));
+
+    policy::SetDMTokenForTesting(
+        policy::DMToken::CreateValidTokenForTesting("ABCDEF"));
+  }
+
+  const Priority priority_ = Priority::FAST_BATCH;
+  const Destination destination_ = Destination::UPLOAD_EVENTS;
+  reporting::MockReportQueue* mock_report_queue_;
+
+  content::BrowserTaskEnvironment task_environment_{
+      base::test::TaskEnvironment::TimeSource::SYSTEM_TIME};
+};
+
+TEST_F(ReportQueueManualTestContextTest,
+       BuildsReportQueueManualTestContextAndUploadsMessages) {
+  uint64_t kNumberOfMessagesToEnqueue = 5;
+  base::TimeDelta kFrequency = base::TimeDelta::FromSeconds(1);
+
+  TestCallbackWaiterWithCounter enqueue_waiter{kNumberOfMessagesToEnqueue};
+  EXPECT_CALL(*mock_report_queue_, StringPieceEnqueue_(_, _, _))
+      .WillRepeatedly(WithArgs<2>(Invoke(
+          [&enqueue_waiter](ReportQueue::EnqueueCallback enqueue_callback) {
+            std::move(enqueue_callback).Run(Status::StatusOK());
+            enqueue_waiter.Signal();
+          })));
+
+  TestCallbackWaiter completion_waiter;
+  auto completion_cb = base::BindOnce(
+      [](TestCallbackWaiter* completion_waiter, Status status) {
+        EXPECT_TRUE(status.ok());
+        completion_waiter->Signal();
+      },
+      &completion_waiter);
+
+  Start<ReportQueueManualTestContext>(
+      kFrequency, kNumberOfMessagesToEnqueue, destination_, priority_,
+      std::move(completion_cb),
+      base::ThreadPool::CreateSequencedTaskRunner(base::TaskTraits()));
+
+  enqueue_waiter.Wait();
+  completion_waiter.Wait();
+}
+
+}  // namespace
+}  // namespace reporting

chrome/test/BUILD.gn

@@ -3647,6 +3647,7 @@ test("unit_tests") {
     "../browser/policy/messaging_layer/upload/record_handler_impl_unittest.cc",
     "../browser/policy/messaging_layer/upload/record_upload_request_builder_unittest.cc",
     "../browser/policy/messaging_layer/upload/upload_client_unittest.cc",
+    "../browser/policy/messaging_layer/util/report_queue_manual_test_context_unittest.cc",
     "../browser/policy/messaging_layer/util/shared_queue_unittest.cc",
     "../browser/policy/messaging_layer/util/shared_vector_unittest.cc",
     "../browser/policy/messaging_layer/util/status_macros_unittest.cc",