Extended the ProtoDatabase to provide LoadKeys() functionality.

components/leveldb_proto/proto_database_impl_unittest.cc

 // Copyright 2014 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 "components/leveldb_proto/proto_database_impl.h"
 
 #include <stddef.h>
 
 #include <map>
 #include <memory>
 #include <utility>
 
 #include "base/bind.h"
 #include "base/files/file_util.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/location.h"
 #include "base/memory/ptr_util.h"
 #include "base/run_loop.h"
 #include "base/threading/thread.h"
 #include "components/leveldb_proto/leveldb_database.h"
 #include "components/leveldb_proto/testing/proto/test.pb.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/leveldatabase/src/include/leveldb/options.h"
 
 using base::MessageLoop;
 using base::ScopedTempDir;
 using testing::Invoke;
 using testing::Return;
+using testing::UnorderedElementsAre;
 using testing::_;
 
 namespace leveldb_proto {
 
 namespace {
 
 typedef std::map<std::string, TestProto> EntryMap;
 
 const char kTestLevelDBClientName[] = "Test";
 
@@ skipping 199 matching lines @@
   ASSERT_TRUE(model.count(key));
   EXPECT_CALL(*mock_db, Get(key, _, _)).WillOnce(SetGetEntry(model));
   EXPECT_CALL(caller, GetCallback1(true, _))
       .WillOnce(VerifyGetEntry(model[key]));
   db_->GetEntry(key, base::Bind(&MockDatabaseCaller::GetCallback,
                                 base::Unretained(&caller)));
 
   base::RunLoop().RunUntilIdle();
 }
 
+TEST(ProtoDatabaseImplLevelDBTest, TestDBSaveAndLoadKeys) {
+  base::MessageLoop main_loop;
+
+  ScopedTempDir temp_dir;
+  ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
+  base::Thread db_thread("dbthread");
+  ASSERT_TRUE(db_thread.Start());
+  std::unique_ptr<ProtoDatabaseImpl<TestProto>> db(
+      new ProtoDatabaseImpl<TestProto>(db_thread.task_runner()));
+
+  auto expect_init_success =
+      base::Bind([](bool success) { EXPECT_TRUE(success); });
+  db->Init(kTestLevelDBClientName, temp_dir.path(), expect_init_success);
+
+  base::RunLoop run_update_entries;
+  auto expect_update_success = base::Bind(
+      [](base::Closure signal, bool success) {
+        EXPECT_TRUE(success);
+        signal.Run();
+      },
+      run_update_entries.QuitClosure());
+  TestProto test_proto;
+  test_proto.set_data("some data");
+  ProtoDatabase<TestProto>::KeyEntryVector data_set(
+          {{"0", test_proto}, {"1", test_proto}, {"2", test_proto}});
+  db->UpdateEntries(
+      base::MakeUnique<ProtoDatabase<TestProto>::KeyEntryVector>(data_set),
+      base::MakeUnique<std::vector<std::string>>(), expect_update_success);
+  run_update_entries.Run();
+
+  base::RunLoop run_load_keys;
+  auto verify_loaded_keys = base::Bind(
+      [](base::Closure signal, bool success,
+         std::unique_ptr<std::vector<std::string>> keys) {
+        EXPECT_TRUE(success);
+        EXPECT_THAT(*keys, UnorderedElementsAre("0", "1", "2"));
+        signal.Run();
+      },
+      run_load_keys.QuitClosure());
+  db->LoadKeys(verify_loaded_keys);
+  run_load_keys.Run();
+
+  // Shutdown database.
+  db.reset();
+  base::RunLoop run_destruction;
+  db_thread.task_runner()->PostTaskAndReply(
+      FROM_HERE, base::Bind(base::DoNothing), run_destruction.QuitClosure());
+  run_destruction.Run();
+}
+
 TEST_F(ProtoDatabaseImplTest, TestDBGetNotFound) {
   base::FilePath path(FILE_PATH_LITERAL("/fake/path"));
 
   MockDB* mock_db = new MockDB();
   MockDatabaseCaller caller;
   EntryMap model = GetSmallModel();
 
   EXPECT_CALL(*mock_db, Init(_));
   EXPECT_CALL(caller, InitCallback(_));
   db_->InitWithDatabase(
@@ skipping 304 matching lines @@
 
   ASSERT_TRUE(db->Save(save_entries, remove_keys));
 
   std::vector<std::string> second_load_entries;
 
   ASSERT_TRUE(db->Load(&second_load_entries));
   EXPECT_EQ(1u, second_load_entries.size());
 }
 
 }  // namespace leveldb_proto