Extended the ProtoDatabase to provide LoadKeys() functionality.

components/ntp_snippets/remote/ntp_snippets_database_unittest.cc

 // Copyright 2016 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/ntp_snippets/remote/ntp_snippets_database.h"
 
 #include <memory>
 
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/files/file_path.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/macros.h"
+#include "base/memory/ptr_util.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 using testing::ElementsAre;
+using testing::Eq;
 using testing::IsEmpty;
 using testing::Mock;
 using testing::_;
 
 namespace ntp_snippets {
 
 bool operator==(const SnippetSource& lhs, const SnippetSource& rhs) {
   return lhs.url == rhs.url && lhs.publisher_name == rhs.publisher_name &&
          lhs.amp_url == rhs.amp_url;
 }
@@ skipping 42 matching lines @@
     // Explicitly destroy any existing database first, so it releases the lock
     // on the file.
     db_.reset();
 
     db_.reset(new NTPSnippetsDatabase(database_dir_.GetPath(),
                                       base::ThreadTaskRunnerHandle::Get()));
   }
 
   NTPSnippetsDatabase* db() { return db_.get(); }
 
+  // TODO(tschumann): MOCK_METHODS on non mock objects are an anti-pattern.
+  // Clean up.
   void OnSnippetsLoaded(NTPSnippet::PtrVector snippets) {
     OnSnippetsLoadedImpl(snippets);
   }
   MOCK_METHOD1(OnSnippetsLoadedImpl,
                void(const NTPSnippet::PtrVector& snippets));
 
   MOCK_METHOD1(OnImageLoaded, void(std::string));
 
  private:
   base::MessageLoop message_loop_;
@@ skipping 218 matching lines @@
   db()->DeleteImage(snippet->id());
 
   // Make sure the image is gone.
   EXPECT_CALL(*this, OnImageLoaded(std::string()));
   db()->LoadImage(snippet->id(),
                   base::Bind(&NTPSnippetsDatabaseTest::OnImageLoaded,
                              base::Unretained(this)));
   base::RunLoop().RunUntilIdle();
 }
 
+void LoadExpectedImage(NTPSnippetsDatabase* db,
+                       const std::string& id,
+                       const std::string& expected_data) {
+  base::RunLoop run_loop;
+  db->LoadImage(id, base::Bind(
+                        [](base::Closure signal, std::string expected_data,
+                           std::string actual_data) {
+                          EXPECT_THAT(actual_data, Eq(expected_data));
+                          signal.Run();
+                        },
+                        run_loop.QuitClosure(), expected_data));
+  run_loop.Run();
+}
+
+TEST_F(NTPSnippetsDatabaseTest, ShouldGarbageCollectImages) {
+  CreateDatabase();
+  base::RunLoop().RunUntilIdle();
+  ASSERT_TRUE(db()->IsInitialized());
+
+  // Store images.
+  db()->SaveImage("snippet-id-1", "pretty-image-1");
+  db()->SaveImage("snippet-id-2", "pretty-image-2");
+  db()->SaveImage("snippet-id-3", "pretty-image-3");
+  base::RunLoop().RunUntilIdle();
+
+  // Make sure the to-be-garbage collected images are there.
+  LoadExpectedImage(db(), "snippet-id-1", "pretty-image-1");
+  LoadExpectedImage(db(), "snippet-id-3", "pretty-image-3");
+
+  // Garbage collect all except the second.
+  db()->GarbageCollectImages(base::MakeUnique<std::set<std::string>>(
+      std::set<std::string>({"snippet-id-2"})));
+  base::RunLoop().RunUntilIdle();
+
+  // Make sure the images are gone.
+  LoadExpectedImage(db(), "snippet-id-1", "");
+  LoadExpectedImage(db(), "snippet-id-3", "");
+  // Make sure the second still exists.
+  LoadExpectedImage(db(), "snippet-id-2", "pretty-image-2");
+}
+
 }  // namespace ntp_snippets