Refactoring of StorageModule-TestStorageModule

1) Make StorageModule to be instantiated only by factory method
2) Make StorageModule overridable by TestStorageModule
3) Allow TestStorageModule to be constructed directly

In the process I started to adjust tests, and ended up with rather
significant refactoring there too:
4) Use TestEvent instead of callback_ and result_ in the test class;
   it allows to make multiple calls in the same test fixture (not
   happening right now, but will likely be needed in the future).
5) Eliminate specialized test classes and use MOCK_METHOD + ON_CALL +
   optionally WillByDefault to achieve the same result. This makes sure
   that later changes in one test class are not skipped in another -
   we now have only one.

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

chrome/browser/policy/messaging_layer/encryption/encryption_module.h

@@ -15,7 +15,7 @@
 namespace reporting {
 
 // TODO(b/153659559) Temporary EncryptionModule until the real one is ready.
-class EncryptionModule : public base::RefCounted<EncryptionModule> {
+class EncryptionModule : public base::RefCountedThreadSafe<EncryptionModule> {
  public:
   EncryptionModule() = default;
 
@@ -30,7 +30,7 @@ class EncryptionModule : public base::RefCounted<EncryptionModule> {
   virtual ~EncryptionModule() = default;
 
  private:
-  friend base::RefCounted<EncryptionModule>;
+  friend base::RefCountedThreadSafe<EncryptionModule>;
 };
 
 }  // namespace reporting

chrome/browser/policy/messaging_layer/encryption/test_encryption_module.cc

@@ -8,18 +8,18 @@
 
 #include "chrome/browser/policy/messaging_layer/util/statusor.h"
 
+using ::testing::Invoke;
+
 namespace reporting {
 namespace test {
 
-StatusOr<std::string> TestEncryptionModule::EncryptRecord(
-    base::StringPiece record) const {
-  return std::string(record);
+TestEncryptionModule::TestEncryptionModule() {
+  ON_CALL(*this, EncryptRecord)
+      .WillByDefault(
+          Invoke([](base::StringPiece record) { return std::string(record); }));
 }
 
-StatusOr<std::string> AlwaysFailsEncryptionModule::EncryptRecord(
-    base::StringPiece record) const {
-  return Status(error::UNKNOWN, "Failing for tests");
-}
+TestEncryptionModule::~TestEncryptionModule() = default;
 
 }  // namespace test
 }  // namespace reporting

chrome/browser/policy/messaging_layer/encryption/test_encryption_module.h

@@ -9,6 +9,7 @@
 
 #include "chrome/browser/policy/messaging_layer/public/report_queue.h"
 #include "chrome/browser/policy/messaging_layer/util/statusor.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace reporting {
@@ -17,23 +18,15 @@ namespace test {
 // An |EncryptionModule| that does no encryption.
 class TestEncryptionModule : public EncryptionModule {
  public:
-  TestEncryptionModule() = default;
+  TestEncryptionModule();
 
-  StatusOr<std::string> EncryptRecord(base::StringPiece record) const override;
+  MOCK_METHOD(StatusOr<std::string>,
+              EncryptRecord,
+              (base::StringPiece record),
+              (const override));
 
  protected:
-  ~TestEncryptionModule() override = default;
-};
-
-// A |TestEncryptionModule| that always fails on |EncryptRecord| calls.
-class AlwaysFailsEncryptionModule final : public TestEncryptionModule {
- public:
-  AlwaysFailsEncryptionModule() = default;
-
-  StatusOr<std::string> EncryptRecord(base::StringPiece record) const override;
-
- protected:
-  ~AlwaysFailsEncryptionModule() override = default;
+  ~TestEncryptionModule() override;
 };
 
 }  // namespace test

chrome/browser/policy/messaging_layer/public/report_client.cc

@@ -21,8 +21,8 @@ namespace reporting {
 
 using base::MakeRefCounted;
 
-ReportingClient::ReportingClient()
-    : storage_(MakeRefCounted<StorageModule>()),
+ReportingClient::ReportingClient(scoped_refptr<StorageModule> storage)
+    : storage_(std::move(storage)),
       encryption_(MakeRefCounted<EncryptionModule>()) {}
 
 ReportingClient::~ReportingClient() = default;
@@ -43,11 +43,13 @@ StatusOr<ReportingClient*> ReportingClient::GetInstance() {
   return instance->ValueOrDie().get();
 }
 
-// TODO(chromium:1078512) As part of completing the StorageModule and
-// EncryptionModule, this create function will need to be updated to check for
-// successful creation of the StorageModule and EncryptionModule.
+// TODO(chromium:1078512) As part of completing the EncryptionModule,
+// this create function will need to be updated to check for
+// successful creation of the EncryptionModule too.
 StatusOr<std::unique_ptr<ReportingClient>> ReportingClient::Create() {
-  auto client = base::WrapUnique<ReportingClient>(new ReportingClient);
+  ASSIGN_OR_RETURN(scoped_refptr<StorageModule> storage,
+                   StorageModule::Create());
+  auto client = base::WrapUnique<ReportingClient>(new ReportingClient(storage));
   return client;
 }
 

chrome/browser/policy/messaging_layer/public/report_client.h

@@ -47,7 +47,7 @@ class ReportingClient {
       std::unique_ptr<ReportQueueConfiguration> config);
 
  private:
-  ReportingClient();
+  explicit ReportingClient(scoped_refptr<StorageModule> storage);
 
   static StatusOr<ReportingClient*> GetInstance();
 

chrome/browser/policy/messaging_layer/public/report_queue_unittest.cc

@@ -26,79 +26,92 @@
 #include "chrome/browser/policy/messaging_layer/util/statusor.h"
 #include "components/policy/core/common/cloud/dm_token.h"
 #include "components/policy/proto/record_constants.pb.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
+using ::policy::DMToken;
+
+using ::testing::_;
+using ::testing::Invoke;
+using ::testing::Return;
+using ::testing::WithArg;
+
+using ::reporting::test::TestEncryptionModule;
+using ::reporting::test::TestStorageModule;
+
 namespace reporting {
 namespace {
 
-using PolicyCheckCallback =
-    reporting::ReportQueueConfiguration::PolicyCheckCallback;
+// Usage (in tests only):
+//
+//   TestEvent<ResType> e;
+//   ... Do some async work passing e.cb() as a completion callback of
+//       base::OnceCallback<void(ResType* res)> type which also may perform
+//       some other action specified by |done| callback provided by the caller.
+//   ... = e.result();  // Will wait for e.cb() to be called and return the
+//                      // collected result.
+//
+// Or, when the callback is not expected to be invoked:
+//
+//   TestEvent<ResType> e(/*expected_to_complete=*/false);
+//   ... Start work passing e.cb() as a completion callback,
+//       which will not happen.
+//
+template <typename ResType>
+class TestEvent {
+ public:
+  explicit TestEvent(bool expected_to_complete = true)
+      : expected_to_complete_(expected_to_complete),
+        completed_(base::WaitableEvent::ResetPolicy::MANUAL,
+                   base::WaitableEvent::InitialState::NOT_SIGNALED) {}
+  ~TestEvent() {
+    if (expected_to_complete_) {
+      EXPECT_TRUE(completed_.IsSignaled()) << "Not responded";
+    } else {
+      EXPECT_FALSE(completed_.IsSignaled()) << "Responded";
+    }
+  }
+  TestEvent(const TestEvent& other) = delete;
+  TestEvent& operator=(const TestEvent& other) = delete;
+  ResType result() {
+    completed_.Wait();
+    return std::forward<ResType>(result_);
+  }
 
-using base::MakeRefCounted;
-using policy::DMToken;
-using reporting::test::AlwaysFailsEncryptionModule;
-using reporting::test::AlwaysFailsStorageModule;
-using reporting::test::TestEncryptionModule;
-using reporting::test::TestStorageModule;
+  // Completion callback to hand over to the processing method.
+  base::OnceCallback<void(ResType res)> cb() {
+    DCHECK(!completed_.IsSignaled());
+    return base::BindOnce(
+        [](base::WaitableEvent* completed, ResType* result, ResType res) {
+          *result = std::forward<ResType>(res);
+          completed->Signal();
+        },
+        base::Unretained(&completed_), base::Unretained(&result_));
+  }
+
+ private:
+  bool expected_to_complete_;
+  base::WaitableEvent completed_;
+  ResType result_;
+};
 
 // Creates a |ReportQueue| using |TestStorageModule| and |TestEncryptionModule|.
 // Allows access to the storage module for checking stored values.
 class ReportQueueTest : public testing::Test {
- public:
+ protected:
   ReportQueueTest()
-      : completed_(base::WaitableEvent::ResetPolicy::MANUAL,
-                   base::WaitableEvent::InitialState::NOT_SIGNALED),
-        result_(error::INTERNAL, "initialized with non-ok status"),
-        storage_module_(MakeRefCounted<TestStorageModule>()),
-        priority_(Priority::IMMEDIATE),
-        dm_token_(DMToken::CreateValidTokenForTesting("FAKE_DM_TOKEN")),
-        destination_(Destination::UPLOAD_EVENTS),
-        encryption_module_(MakeRefCounted<TestEncryptionModule>()),
-        policy_check_callback_(base::BindRepeating(
-            []() -> Status { return Status::StatusOK(); })) {}
-
-  // Allows specifying an alternative |TestStorageModule| for testing different
-  // cases.
-  explicit ReportQueueTest(scoped_refptr<TestStorageModule> storage_module)
-      : completed_(base::WaitableEvent::ResetPolicy::MANUAL,
-                   base::WaitableEvent::InitialState::NOT_SIGNALED),
-        result_(error::INTERNAL, "initialized with non-ok status"),
-        storage_module_(storage_module),
-        priority_(Priority::IMMEDIATE),
-        dm_token_(DMToken::CreateValidTokenForTesting("FAKE_DM_TOKEN")),
-        destination_(Destination::UPLOAD_EVENTS),
-        encryption_module_(MakeRefCounted<TestEncryptionModule>()),
-        policy_check_callback_(base::BindRepeating(
-            []() -> Status { return Status::StatusOK(); })) {}
-
-  // Allows specifying an alternative |TestEncryptionModule| for testing
-  // different cases.
-  explicit ReportQueueTest(
-      scoped_refptr<TestEncryptionModule> encryption_module)
-      : completed_(base::WaitableEvent::ResetPolicy::MANUAL,
-                   base::WaitableEvent::InitialState::NOT_SIGNALED),
-        result_(error::INTERNAL, "initialized with non-ok status"),
-        storage_module_(MakeRefCounted<TestStorageModule>()),
-        priority_(Priority::IMMEDIATE),
+      : priority_(Priority::IMMEDIATE),
         dm_token_(DMToken::CreateValidTokenForTesting("FAKE_DM_TOKEN")),
         destination_(Destination::UPLOAD_EVENTS),
-        encryption_module_(encryption_module),
-        policy_check_callback_(base::BindRepeating(
-            []() -> Status { return Status::StatusOK(); })) {}
-
-  explicit ReportQueueTest(
-      ReportQueueConfiguration::PolicyCheckCallback policy_check_callback)
-      : completed_(base::WaitableEvent::ResetPolicy::MANUAL,
-                   base::WaitableEvent::InitialState::NOT_SIGNALED),
-        result_(error::INTERNAL, "initialized with non-ok status"),
-        storage_module_(MakeRefCounted<TestStorageModule>()),
-        priority_(Priority::IMMEDIATE),
-        dm_token_(DMToken::CreateValidTokenForTesting("FAKE_DM_TOKEN")),
-        destination_(Destination::UPLOAD_EVENTS),
-        encryption_module_(MakeRefCounted<TestEncryptionModule>()),
-        policy_check_callback_(std::move(policy_check_callback)) {}
+        storage_module_(base::MakeRefCounted<TestStorageModule>()),
+        encryption_module_(base::MakeRefCounted<TestEncryptionModule>()),
+        policy_check_callback_(
+            base::BindRepeating(&ReportQueueTest::MockedPolicyCheck,
+                                base::Unretained(this))) {}
 
   void SetUp() override {
+    ON_CALL(*this, MockedPolicyCheck).WillByDefault(Return(Status::StatusOK()));
+
     StatusOr<std::unique_ptr<ReportQueueConfiguration>> config_result =
         ReportQueueConfiguration::Create(dm_token_, destination_, priority_,
                                          policy_check_callback_);
@@ -112,23 +125,28 @@ class ReportQueueTest : public testing::Test {
     ASSERT_TRUE(report_queue_result.ok());
 
     report_queue_ = std::move(report_queue_result.ValueOrDie());
+  }
 
-    callback_ = base::BindOnce(
-        [](base::WaitableEvent* completed, Status* result,
-           Status status) -> void {
-          *result = status;
-          completed->Signal();
-        },
-        &completed_, &result_);
+  TestStorageModule* test_storage_module() const {
+    TestStorageModule* test_storage_module =
+        google::protobuf::down_cast<TestStorageModule*>(storage_module_.get());
+    DCHECK(test_storage_module);
+    return test_storage_module;
   }
 
- protected:
-  base::WaitableEvent completed_;
-  Status result_;
+  TestEncryptionModule* test_encryption_module() const {
+    TestEncryptionModule* test_encryption_module =
+        google::protobuf::down_cast<TestEncryptionModule*>(
+            encryption_module_.get());
+    DCHECK(test_encryption_module);
+    return test_encryption_module;
+  }
+
+  MOCK_METHOD(Status, MockedPolicyCheck, (), ());
+
   base::test::TaskEnvironment task_envrionment_{
       base::test::TaskEnvironment::TimeSource::MOCK_TIME};
 
-  scoped_refptr<TestStorageModule> storage_module_;
   const Priority priority_;
 
   std::unique_ptr<ReportQueue> report_queue_;
@@ -137,24 +155,24 @@ class ReportQueueTest : public testing::Test {
  private:
   const DMToken dm_token_;
   const Destination destination_;
-  scoped_refptr<TestEncryptionModule> encryption_module_;
-  PolicyCheckCallback policy_check_callback_;
+  scoped_refptr<StorageModule> storage_module_;
+  scoped_refptr<EncryptionModule> encryption_module_;
+  ReportQueueConfiguration::PolicyCheckCallback policy_check_callback_;
 };
 
 // Enqueues a random string and ensures that the string arrives unaltered in the
 // |StorageModule|.
 TEST_F(ReportQueueTest, SuccessfulStringRecord) {
   constexpr char kTestString[] = "El-Chupacabra";
-  Status status = report_queue_->Enqueue(kTestString, std::move(callback_));
-  ASSERT_TRUE(status.ok());
-
-  completed_.Wait();
-
-  EXPECT_TRUE(result_.ok());
+  TestEvent<Status> a;
+  Status status = report_queue_->Enqueue(kTestString, a.cb());
+  ASSERT_OK(status);
+  EXPECT_OK(a.result());
 
-  EXPECT_EQ(storage_module_->priority(), priority_);
+  EXPECT_EQ(test_storage_module()->priority(), priority_);
 
-  EXPECT_EQ(storage_module_->wrapped_record().record().data(), kTestString);
+  EXPECT_EQ(test_storage_module()->wrapped_record().record().data(),
+            kTestString);
 }
 
 // Enqueues a |base::Value| dictionary and ensures it arrives unaltered in the
@@ -164,17 +182,15 @@ TEST_F(ReportQueueTest, SuccessfulBaseValueRecord) {
   constexpr char kTestValue[] = "TEST_VALUE";
   base::Value test_dict(base::Value::Type::DICTIONARY);
   test_dict.SetStringKey(kTestKey, kTestValue);
-  Status status = report_queue_->Enqueue(test_dict, std::move(callback_));
-  ASSERT_TRUE(status.ok());
+  TestEvent<Status> a;
+  Status status = report_queue_->Enqueue(test_dict, a.cb());
+  ASSERT_OK(status);
+  EXPECT_OK(a.result());
 
-  completed_.Wait();
+  EXPECT_EQ(test_storage_module()->priority(), priority_);
 
-  EXPECT_TRUE(result_.ok());
-
-  EXPECT_EQ(storage_module_->priority(), priority_);
-
-  base::Optional<base::Value> value_result =
-      base::JSONReader::Read(storage_module_->wrapped_record().record().data());
+  base::Optional<base::Value> value_result = base::JSONReader::Read(
+      test_storage_module()->wrapped_record().record().data());
   ASSERT_TRUE(value_result);
   EXPECT_EQ(value_result.value(), test_dict);
 }
@@ -184,94 +200,85 @@ TEST_F(ReportQueueTest, SuccessfulBaseValueRecord) {
 TEST_F(ReportQueueTest, SuccessfulProtoRecord) {
   reporting::test::TestMessage test_message;
   test_message.set_test("TEST_MESSAGE");
-  Status status = report_queue_->Enqueue(&test_message, std::move(callback_));
-  ASSERT_TRUE(status.ok());
-
-  completed_.Wait();
+  TestEvent<Status> a;
+  Status status = report_queue_->Enqueue(&test_message, a.cb());
+  ASSERT_OK(status);
+  EXPECT_OK(a.result());
 
-  EXPECT_TRUE(result_.ok());
-
-  EXPECT_EQ(storage_module_->priority(), priority_);
+  EXPECT_EQ(test_storage_module()->priority(), priority_);
 
   reporting::test::TestMessage result_message;
   ASSERT_TRUE(result_message.ParseFromString(
-      storage_module_->wrapped_record().record().data()));
+      test_storage_module()->wrapped_record().record().data()));
   ASSERT_EQ(result_message.test(), test_message.test());
 }
 
-// A |ReportQueueTest| built with an |AlwaysFailsStorageModule|.
-class StorageFailsReportQueueTest : public ReportQueueTest {
- public:
-  StorageFailsReportQueueTest()
-      : ReportQueueTest(MakeRefCounted<AlwaysFailsStorageModule>()) {}
-};
-
 // The call to enqueue should succeed, indicating that the storage operation has
 // been scheduled. The callback should fail, indicating that storage was
 // unsuccessful.
-TEST_F(StorageFailsReportQueueTest, CallSuccessCallbackFailure) {
+TEST_F(ReportQueueTest, CallSuccessCallbackFailure) {
+  EXPECT_CALL(*test_storage_module(), AddRecord(_, _, _))
+      .WillOnce(
+          WithArg<2>(Invoke([](base::OnceCallback<void(Status)> callback) {
+            std::move(callback).Run(Status(error::UNKNOWN, "Failing for Test"));
+          })));
+
   reporting::test::TestMessage test_message;
   test_message.set_test("TEST_MESSAGE");
-  Status status = report_queue_->Enqueue(&test_message, std::move(callback_));
-  ASSERT_TRUE(status.ok());
-
-  completed_.Wait();
-
-  EXPECT_FALSE(result_.ok());
-  EXPECT_EQ(result_.error_code(), error::UNKNOWN);
+  TestEvent<Status> a;
+  Status status = report_queue_->Enqueue(&test_message, a.cb());
+  ASSERT_OK(status);
+  auto result = a.result();
+  EXPECT_FALSE(result.ok());
+  EXPECT_EQ(result.error_code(), error::UNKNOWN);
 }
 
-// A |ReportQueueTest| built with an |AlwaysFailsEncryptionModule|.
-class EncryptionFailsReportQueueTest : public ReportQueueTest {
- public:
-  EncryptionFailsReportQueueTest()
-      : ReportQueueTest(MakeRefCounted<AlwaysFailsEncryptionModule>()) {}
-};
-
 // The call to enqueue should succeed, indicating that the encryption operation
 // has been scheduled. The callback should fail, indicating that encryption was
 // unsuccessful.
-TEST_F(EncryptionFailsReportQueueTest, CallSuccessCallFailure) {
+TEST_F(ReportQueueTest, EnqueueSuccessEncryptFailure) {
+  EXPECT_CALL(*test_encryption_module(), EncryptRecord(_))
+      .WillOnce(Return(Status(error::UNKNOWN, "Failing for tests")));
   reporting::test::TestMessage test_message;
   test_message.set_test("TEST_MESSAGE");
-  Status status = report_queue_->Enqueue(&test_message, std::move(callback_));
-  ASSERT_TRUE(status.ok());
-
-  completed_.Wait();
-
-  EXPECT_FALSE(result_.ok());
-  EXPECT_EQ(result_.error_code(), error::UNKNOWN);
+  TestEvent<Status> a;
+  Status status = report_queue_->Enqueue(&test_message, a.cb());
+  ASSERT_OK(status);
+  auto result = a.result();
+  EXPECT_FALSE(result.ok());
+  EXPECT_EQ(result.error_code(), error::UNKNOWN);
 }
 
-class ReportQueueRespectsPolicyTest : public ReportQueueTest {
- public:
-  ReportQueueRespectsPolicyTest()
-      : ReportQueueTest(base::BindRepeating([]() -> Status {
-          return Status(error::UNAUTHENTICATED, "Failing for tests");
-        })) {}
-};
-
-TEST_F(ReportQueueRespectsPolicyTest, EnqueueStringFailsOnPolicy) {
+TEST_F(ReportQueueTest, EnqueueStringFailsOnPolicy) {
+  EXPECT_CALL(*this, MockedPolicyCheck)
+      .WillOnce(Return(Status(error::UNAUTHENTICATED, "Failing for tests")));
   constexpr char kTestString[] = "El-Chupacabra";
-  Status status = report_queue_->Enqueue(kTestString, std::move(callback_));
+  TestEvent<Status> a(/*expected_to_complete=*/false);
+  Status status = report_queue_->Enqueue(kTestString, a.cb());
   EXPECT_FALSE(status.ok());
   EXPECT_EQ(status.error_code(), error::UNAUTHENTICATED);
 }
 
-TEST_F(ReportQueueRespectsPolicyTest, EnqueueProtoFailsOnPolicy) {
+TEST_F(ReportQueueTest, EnqueueProtoFailsOnPolicy) {
+  EXPECT_CALL(*this, MockedPolicyCheck)
+      .WillOnce(Return(Status(error::UNAUTHENTICATED, "Failing for tests")));
   reporting::test::TestMessage test_message;
   test_message.set_test("TEST_MESSAGE");
-  Status status = report_queue_->Enqueue(&test_message, std::move(callback_));
+  TestEvent<Status> a(/*expected_to_complete=*/false);
+  Status status = report_queue_->Enqueue(&test_message, a.cb());
   EXPECT_FALSE(status.ok());
   EXPECT_EQ(status.error_code(), error::UNAUTHENTICATED);
 }
 
-TEST_F(ReportQueueRespectsPolicyTest, EnqueueValueFailsOnPolicy) {
+TEST_F(ReportQueueTest, EnqueueValueFailsOnPolicy) {
+  EXPECT_CALL(*this, MockedPolicyCheck)
+      .WillOnce(Return(Status(error::UNAUTHENTICATED, "Failing for tests")));
   constexpr char kTestKey[] = "TEST_KEY";
   constexpr char kTestValue[] = "TEST_VALUE";
   base::Value test_dict(base::Value::Type::DICTIONARY);
   test_dict.SetStringKey(kTestKey, kTestValue);
-  Status status = report_queue_->Enqueue(test_dict, std::move(callback_));
+  TestEvent<Status> a(/*expected_to_complete=*/false);
+  Status status = report_queue_->Enqueue(test_dict, a.cb());
   EXPECT_FALSE(status.ok());
   EXPECT_EQ(status.error_code(), error::UNAUTHENTICATED);
 }

chrome/browser/policy/messaging_layer/storage/storage.h

@@ -45,6 +45,7 @@ class Storage : public base::RefCountedThreadSafe<Storage> {
     virtual void ProcessBlob(Priority priority,
                              StatusOr<base::span<const uint8_t>> data,
                              base::OnceCallback<void(bool)> processed_cb) = 0;
+
     // Finalizes the upload (e.g. sends the message to the server and gets
     // response).
     virtual void Completed(Priority priority, Status final_status) = 0;

chrome/browser/policy/messaging_layer/storage/storage_module.cc

@@ -5,13 +5,19 @@
 #include <utility>
 
 #include "base/callback.h"
+#include "base/no_destructor.h"
 #include "chrome/browser/policy/messaging_layer/storage/storage_module.h"
 #include "chrome/browser/policy/messaging_layer/util/status.h"
+#include "chrome/browser/policy/messaging_layer/util/statusor.h"
 #include "components/policy/proto/record.pb.h"
 #include "components/policy/proto/record_constants.pb.h"
 
 namespace reporting {
 
+StorageModule::StorageModule() = default;
+
+StorageModule::~StorageModule() = default;
+
 void StorageModule::AddRecord(reporting::EncryptedRecord record,
                               reporting::Priority priority,
                               base::OnceCallback<void(Status)> callback) {
@@ -19,4 +25,12 @@ void StorageModule::AddRecord(reporting::EncryptedRecord record,
       Status(error::UNIMPLEMENTED, "AddRecord isn't implemented"));
 }
 
+// static
+StatusOr<scoped_refptr<StorageModule>> StorageModule::Create() {
+  scoped_refptr<StorageModule> instance =
+      // Cannot base::MakeRefCounted, since constructor is protected.
+      base::WrapRefCounted(new StorageModule());
+  return instance;
+}
+
 }  // namespace reporting

chrome/browser/policy/messaging_layer/storage/storage_module.h

@@ -10,15 +10,17 @@
 #include "base/callback.h"
 #include "base/memory/ref_counted.h"
 #include "chrome/browser/policy/messaging_layer/util/status.h"
+#include "chrome/browser/policy/messaging_layer/util/statusor.h"
 #include "components/policy/proto/record.pb.h"
 #include "components/policy/proto/record_constants.pb.h"
 
 namespace reporting {
 
 // TODO(b/153659559) Temporary StorageModule until the real one is ready.
-class StorageModule : public base::RefCounted<StorageModule> {
+class StorageModule : public base::RefCountedThreadSafe<StorageModule> {
  public:
-  StorageModule() = default;
+  // Factory method creates |StorageModule| object.
+  static StatusOr<scoped_refptr<StorageModule>> Create();
 
   StorageModule(const StorageModule& other) = delete;
   StorageModule& operator=(const StorageModule& other) = delete;
@@ -30,10 +32,14 @@ class StorageModule : public base::RefCounted<StorageModule> {
                          base::OnceCallback<void(Status)> callback);
 
  protected:
-  virtual ~StorageModule() = default;
+  // Constructor can only be called by |Create| factory method.
+  StorageModule();
+
+  // Refcounted object must have destructor declared protected or private.
+  virtual ~StorageModule();
 
  private:
-  friend base::RefCounted<StorageModule>;
+  friend base::RefCountedThreadSafe<StorageModule>;
 };
 
 }  // namespace reporting

chrome/browser/policy/messaging_layer/storage/test_storage_module.cc

@@ -10,29 +10,41 @@
 #include "chrome/browser/policy/messaging_layer/public/report_queue.h"
 #include "components/policy/proto/record.pb.h"
 #include "components/policy/proto/record_constants.pb.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
+using ::testing::Invoke;
+
 namespace reporting {
 namespace test {
 
-using reporting::EncryptedRecord;
-using reporting::Priority;
+TestStorageModule::TestStorageModule() {
+  ON_CALL(*this, AddRecord)
+      .WillByDefault(Invoke(this, &TestStorageModule::AddRecordSuccessfully));
+}
 
-void TestStorageModule::AddRecord(EncryptedRecord record,
-                                  Priority priority,
-                                  base::OnceCallback<void(Status)> callback) {
-  ASSERT_TRUE(
-      wrapped_record_.ParseFromString(record.encrypted_wrapped_record()));
-  priority_ = priority;
-  std::move(callback).Run(Status::StatusOK());
+TestStorageModule::~TestStorageModule() = default;
+
+WrappedRecord TestStorageModule::wrapped_record() const {
+  EXPECT_TRUE(wrapped_record_.has_value());
+  return wrapped_record_.value();
 }
 
-void AlwaysFailsStorageModule::AddRecord(
+Priority TestStorageModule::priority() const {
+  EXPECT_TRUE(priority_.has_value());
+  return priority_.value();
+}
+
+void TestStorageModule::AddRecordSuccessfully(
     EncryptedRecord record,
     Priority priority,
     base::OnceCallback<void(Status)> callback) {
-  std::move(callback).Run(Status(error::UNKNOWN, "Failing for Tests"));
+  WrappedRecord wrapped_record;
+  ASSERT_TRUE(
+      wrapped_record.ParseFromString(record.encrypted_wrapped_record()));
+  wrapped_record_ = wrapped_record;
+  priority_ = priority;
+  std::move(callback).Run(Status::StatusOK());
 }
-
 }  // namespace test
 }  // namespace reporting

chrome/browser/policy/messaging_layer/storage/test_storage_module.h

@@ -8,46 +8,40 @@
 #include <utility>
 
 #include "base/callback.h"
+#include "base/optional.h"
 #include "chrome/browser/policy/messaging_layer/public/report_queue.h"
 #include "components/policy/proto/record.pb.h"
 #include "components/policy/proto/record_constants.pb.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
 
 namespace reporting {
 namespace test {
 
-// A |StorageModule| that stores the wrapped record and priority and calls the
-// callback with an OK status.
 class TestStorageModule : public StorageModule {
  public:
-  TestStorageModule() = default;
+  TestStorageModule();
 
-  void AddRecord(reporting::EncryptedRecord record,
-                 reporting::Priority priority,
-                 base::OnceCallback<void(Status)> callback) override;
+  MOCK_METHOD(void,
+              AddRecord,
+              (EncryptedRecord record,
+               Priority priority,
+               base::OnceCallback<void(Status)> callback),
+              (override));
 
-  reporting::WrappedRecord wrapped_record() { return wrapped_record_; }
-
-  reporting::Priority priority() { return priority_; }
+  WrappedRecord wrapped_record() const;
+  Priority priority() const;
 
  protected:
-  ~TestStorageModule() override = default;
+  ~TestStorageModule() override;
 
  private:
-  reporting::WrappedRecord wrapped_record_;
-  reporting::Priority priority_;
-};
-
-// A |TestStorageModule| that always fails on |AddRecord| calls.
-class AlwaysFailsStorageModule final : public TestStorageModule {
- public:
-  AlwaysFailsStorageModule() = default;
-
-  void AddRecord(reporting::EncryptedRecord record,
-                 reporting::Priority priority,
-                 base::OnceCallback<void(Status)> callback) override;
+  void AddRecordSuccessfully(EncryptedRecord record,
+                             Priority priority,
+                             base::OnceCallback<void(Status)> callback);
 
- protected:
-  ~AlwaysFailsStorageModule() override = default;
+  base::Optional<WrappedRecord> wrapped_record_;
+  base::Optional<Priority> priority_;
 };
 
 }  // namespace test