[NTP::SectionOrder] Replace CategoryFactory with a category ranker.

components/ntp_snippets/remote/remote_suggestions_provider_unittest.cc

 // Copyright 2015 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/remote_suggestions_provider.h"
 
 #include <memory>
 #include <utility>
 #include <vector>
 
 #include "base/command_line.h"
 #include "base/files/file_path.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/json/json_reader.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/strings/string_number_conversions.h"
 #include "base/strings/string_util.h"
 #include "base/strings/stringprintf.h"
 #include "base/test/histogram_tester.h"
 #include "base/test/simple_test_clock.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/time.h"
 #include "components/image_fetcher/image_decoder.h"
 #include "components/image_fetcher/image_fetcher.h"
 #include "components/image_fetcher/image_fetcher_delegate.h"
-#include "components/ntp_snippets/category_factory.h"
+#include "components/ntp_snippets/category.h"
 #include "components/ntp_snippets/category_info.h"
+#include "components/ntp_snippets/category_rankers/category_ranker.h"
+#include "components/ntp_snippets/category_rankers/constant_category_ranker.h"
+#include "components/ntp_snippets/category_rankers/mock_category_ranker.h"
 #include "components/ntp_snippets/ntp_snippets_constants.h"
 #include "components/ntp_snippets/pref_names.h"
 #include "components/ntp_snippets/remote/ntp_snippet.h"
 #include "components/ntp_snippets/remote/ntp_snippets_fetcher.h"
 #include "components/ntp_snippets/remote/ntp_snippets_scheduler.h"
 #include "components/ntp_snippets/remote/remote_suggestions_database.h"
 #include "components/ntp_snippets/remote/test_utils.h"
 #include "components/ntp_snippets/user_classifier.h"
 #include "components/prefs/testing_pref_service.h"
 #include "components/signin/core/browser/fake_profile_oauth2_token_service.h"
 #include "components/signin/core/browser/fake_signin_manager.h"
 #include "components/variations/variations_params_manager.h"
 #include "net/url_request/test_url_fetcher_factory.h"
 #include "net/url_request/url_request_test_util.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "ui/gfx/image/image.h"
 #include "ui/gfx/image/image_unittest_util.h"
 
 using image_fetcher::ImageFetcher;
 using image_fetcher::ImageFetcherDelegate;
+using testing::_;
 using testing::ElementsAre;
 using testing::Eq;
 using testing::InSequence;
 using testing::Invoke;
 using testing::IsEmpty;
 using testing::Mock;
 using testing::MockFunction;
 using testing::NiceMock;
 using testing::Not;
 using testing::SaveArg;
 using testing::SizeIs;
 using testing::StartsWith;
 using testing::WithArgs;
-using testing::_;
 
 namespace ntp_snippets {
 
 namespace {
 
 MATCHER_P(IdEq, value, "") {
   return arg->id() == value;
 }
 
 MATCHER_P(IdWithinCategoryEq, expected_id, "") {
@@ skipping 33 matching lines @@
 base::Time GetDefaultCreationTime() {
   base::Time out_time;
   EXPECT_TRUE(base::Time::FromUTCExploded(kDefaultCreationTime, &out_time));
   return out_time;
 }
 
 base::Time GetDefaultExpirationTime() {
   return base::Time::Now() + base::TimeDelta::FromHours(1);
 }
 
+std::string GetCategoryJson(const std::vector<std::string>& snippets,
+                            int remote_category_id,
+                            const std::string& category_title) {
+  return base::StringPrintf(
+      "  {\n"
+      "    \"id\": %d,\n"
+      "    \"localizedTitle\": \"%s\",\n"
+      "    \"suggestions\": [%s]\n"
+      "  }\n",
+      remote_category_id, category_title.c_str(),
+      base::JoinString(snippets, ", ").c_str());
+}
+
+class MultiCategoryJsonBuilder {
+ public:
+  MultiCategoryJsonBuilder() {}
+
+  MultiCategoryJsonBuilder& AddCategoryWithCustomTitle(
+      const std::vector<std::string>& snippets,
+      int remote_category_id,
+      const std::string& category_title) {
+    category_json_.push_back(
+        GetCategoryJson(snippets, remote_category_id, category_title));
+    return *this;
+  }
+
+  MultiCategoryJsonBuilder& AddCategory(
+      const std::vector<std::string>& snippets,
+      int remote_category_id) {
+    return AddCategoryWithCustomTitle(
+        snippets, remote_category_id,
+        "Title" + base::IntToString(remote_category_id));
+  }
+
+  std::string Build() {
+    return base::StringPrintf(
+        "{\n"
+        "  \"categories\": [\n"
+        "%s\n"
+        "  ]\n"
+        "}\n",
+        base::JoinString(category_json_, "  ,\n").c_str());
+  }
+
+ private:
+  std::vector<std::string> category_json_;
+};
+
+// TODO(vitaliii): Remove these convenience functions as they do not provide
+// that much value and add additional redirections obscuring the code.
 std::string GetTestJson(const std::vector<std::string>& snippets,
                         const std::string& category_title) {
-  return base::StringPrintf(
-      "{\n"
-      "  \"categories\": [{\n"
-      "    \"id\": 1,\n"
-      "    \"localizedTitle\": \"%s\",\n"
-      "    \"suggestions\": [%s]\n"
-      "  }]\n"
-      "}\n",
-      category_title.c_str(), base::JoinString(snippets, ", ").c_str());
+  return MultiCategoryJsonBuilder()
+      .AddCategoryWithCustomTitle(snippets, /*remote_category_id=*/1,
+                                  category_title)
+      .Build();
 }
 
 std::string GetTestJson(const std::vector<std::string>& snippets) {
   return GetTestJson(snippets, kTestJsonDefaultCategoryTitle);
 }
 
-// TODO(tschumann): Remove the default parameter other_id. It makes the tests
-// less explicit and hard to read. Also get rid of the convenience
-// other_category() and unknown_category() helpers -- tests can just define
-// their own.
-std::string GetMultiCategoryJson(const std::vector<std::string>& articles,
-                                 const std::vector<std::string>& others,
-                                 int other_id = 2) {
-  return base::StringPrintf(
-      "{\n"
-      "  \"categories\": [{\n"
-      "    \"id\": 1,\n"
-      "    \"localizedTitle\": \"Articles for You\",\n"
-      "    \"suggestions\": [%s]\n"
-      "  }, {\n"
-      "    \"id\": %i,\n"
-      "    \"localizedTitle\": \"Other Things\",\n"
-      "    \"suggestions\": [%s]\n"
-      "  }]\n"
-      "}\n",
-      base::JoinString(articles, ", ").c_str(), other_id,
-      base::JoinString(others, ", ").c_str());
-}
-
 std::string FormatTime(const base::Time& t) {
   base::Time::Exploded x;
   t.UTCExplode(&x);
   return base::StringPrintf("%04d-%02d-%02dT%02d:%02d:%02dZ", x.year, x.month,
                             x.day_of_month, x.hour, x.minute, x.second);
 }
 
 std::string GetSnippetWithUrlAndTimesAndSource(
     const std::vector<std::string>& ids,
     const std::string& url,
@@ skipping 225 matching lines @@
 class RemoteSuggestionsProviderTest : public ::testing::Test {
  public:
   RemoteSuggestionsProviderTest()
       : params_manager_(ntp_snippets::kStudyName,
                         {{"content_suggestions_backend",
                           kTestContentSuggestionsServerEndpoint},
                          {"fetching_personalization", "non_personal"}}),
         fake_url_fetcher_factory_(
             /*default_factory=*/&failing_url_fetcher_factory_),
         test_url_(kTestContentSuggestionsServerWithAPIKey),
+        category_ranker_(base::MakeUnique<ConstantCategoryRanker>()),
         user_classifier_(/*pref_service=*/nullptr),
         image_fetcher_(nullptr),
         image_decoder_(nullptr),
         database_(nullptr) {
     RemoteSuggestionsProvider::RegisterProfilePrefs(
         utils_.pref_service()->registry());
     RequestThrottler::RegisterProfilePrefs(utils_.pref_service()->registry());
 
     EXPECT_TRUE(database_dir_.CreateUniqueTempDir());
   }
 
   ~RemoteSuggestionsProviderTest() override {
     // We need to run the message loop after deleting the database, because
     // ProtoDatabaseImpl deletes the actual LevelDB asynchronously on the task
     // runner. Without this, we'd get reports of memory leaks.
     base::RunLoop().RunUntilIdle();
   }
 
+  // TODO(vitaliii): Rewrite this function to initialize a test class member
+  // instead of creating a new service.
   std::unique_ptr<RemoteSuggestionsProvider> MakeSnippetsService(
       bool set_empty_response = true) {
     auto service = MakeSnippetsServiceWithoutInitialization();
     WaitForSnippetsServiceInitialization(service.get(), set_empty_response);
     return service;
   }
 
   std::unique_ptr<RemoteSuggestionsProvider>
   MakeSnippetsServiceWithoutInitialization() {
     scoped_refptr<base::SingleThreadTaskRunner> task_runner(
         base::ThreadTaskRunnerHandle::Get());
     scoped_refptr<net::TestURLRequestContextGetter> request_context_getter =
         new net::TestURLRequestContextGetter(task_runner.get());
 
     utils_.ResetSigninManager();
     std::unique_ptr<NTPSnippetsFetcher> snippets_fetcher =
         base::MakeUnique<NTPSnippetsFetcher>(
             utils_.fake_signin_manager(), fake_token_service_.get(),
-            std::move(request_context_getter), utils_.pref_service(),
-            &category_factory_, nullptr, base::Bind(&ParseJson), kAPIKey,
-            &user_classifier_);
+            std::move(request_context_getter), utils_.pref_service(), nullptr,
+            base::Bind(&ParseJson), kAPIKey, &user_classifier_);
 
     utils_.fake_signin_manager()->SignIn("foo@bar.com");
 
     auto image_fetcher = base::MakeUnique<NiceMock<MockImageFetcher>>();
 
     image_fetcher_ = image_fetcher.get();
     EXPECT_CALL(*image_fetcher, SetImageFetcherDelegate(_));
     auto image_decoder = base::MakeUnique<FakeImageDecoder>();
     image_decoder_ = image_decoder.get();
     EXPECT_FALSE(observer_);
     observer_ = base::MakeUnique<FakeContentSuggestionsProviderObserver>();
     auto database = base::MakeUnique<RemoteSuggestionsDatabase>(
         database_dir_.GetPath(), task_runner);
     database_ = database.get();
     return base::MakeUnique<RemoteSuggestionsProvider>(
-        observer_.get(), &category_factory_, utils_.pref_service(), "fr",
+        observer_.get(), utils_.pref_service(), "fr", category_ranker_.get(),
         &user_classifier_, &scheduler_, std::move(snippets_fetcher),
-        std::move(image_fetcher), std::move(image_decoder),
-        std::move(database),
+        std::move(image_fetcher), std::move(image_decoder), std::move(database),
         base::MakeUnique<RemoteSuggestionsStatusService>(
             utils_.fake_signin_manager(), utils_.pref_service()));
   }
 
   void WaitForSnippetsServiceInitialization(RemoteSuggestionsProvider* service,
                                             bool set_empty_response) {
     EXPECT_EQ(RemoteSuggestionsProvider::State::NOT_INITED, service->state_);
 
     // Add an initial fetch response, as the service tries to fetch when there
     // is nothing in the DB.
     if (set_empty_response) {
       SetUpFetchResponse(GetTestJson(std::vector<std::string>()));
     }
 
     // TODO(treib): Find a better way to wait for initialization to finish.
     base::RunLoop().RunUntilIdle();
     EXPECT_NE(RemoteSuggestionsProvider::State::NOT_INITED, service->state_);
   }
 
-  void ResetSnippetsService(
-      std::unique_ptr<RemoteSuggestionsProvider>* service) {
+  void ResetSnippetsService(std::unique_ptr<RemoteSuggestionsProvider>* service,
+                            bool set_empty_response) {
     service->reset();
     observer_.reset();
-    *service = MakeSnippetsService();
+    *service = MakeSnippetsService(set_empty_response);
+  }
+
+  void SetCategoryRanker(std::unique_ptr<CategoryRanker> category_ranker) {
+    category_ranker_ = std::move(category_ranker);
   }
 
   ContentSuggestion::ID MakeArticleID(const std::string& id_within_category) {
     return ContentSuggestion::ID(articles_category(), id_within_category);
   }
 
   Category articles_category() {
-    return category_factory_.FromKnownCategory(KnownCategories::ARTICLES);
+    return Category::FromKnownCategory(KnownCategories::ARTICLES);
   }
 
   ContentSuggestion::ID MakeOtherID(const std::string& id_within_category) {
     return ContentSuggestion::ID(other_category(), id_within_category);
   }
 
-  Category other_category() { return category_factory_.FromRemoteCategory(2); }
+  // TODO(tschumann): Get rid of the convenience other_category() and
+  // unknown_category() helpers -- tests can just define their own.
+  Category other_category() { return Category::FromRemoteCategory(2); }
 
   Category unknown_category() {
-    return category_factory_.FromRemoteCategory(kUnknownRemoteCategoryId);
+    return Category::FromRemoteCategory(kUnknownRemoteCategoryId);
   }
 
  protected:
   const GURL& test_url() { return test_url_; }
   FakeContentSuggestionsProviderObserver& observer() { return *observer_; }
   MockScheduler& mock_scheduler() { return scheduler_; }
   // TODO(tschumann): Make this a strict-mock. We want to avoid unneccesary
   // network requests.
   NiceMock<MockImageFetcher>* image_fetcher() { return image_fetcher_; }
   FakeImageDecoder* image_decoder() { return image_decoder_; }
@@ skipping 32 matching lines @@
 
  private:
   variations::testing::VariationParamsManager params_manager_;
   test::RemoteSuggestionsTestUtils utils_;
   base::MessageLoop message_loop_;
   FailingFakeURLFetcherFactory failing_url_fetcher_factory_;
   // Instantiation of factory automatically sets itself as URLFetcher's factory.
   net::FakeURLFetcherFactory fake_url_fetcher_factory_;
   const GURL test_url_;
   std::unique_ptr<OAuth2TokenService> fake_token_service_;
+  std::unique_ptr<CategoryRanker> category_ranker_;
   UserClassifier user_classifier_;
   NiceMock<MockScheduler> scheduler_;
   std::unique_ptr<FakeContentSuggestionsProviderObserver> observer_;
-  CategoryFactory category_factory_;
   NiceMock<MockImageFetcher>* image_fetcher_;
   FakeImageDecoder* image_decoder_;
 
   base::ScopedTempDir database_dir_;
   RemoteSuggestionsDatabase* database_;
 
   DISALLOW_COPY_AND_ASSIGN(RemoteSuggestionsProviderTest);
 };
 
 TEST_F(RemoteSuggestionsProviderTest, ScheduleOnStart) {
@@ skipping 14 matching lines @@
   SetUpFetchResponse(GetTestJson({GetSnippet()}));
   auto service = MakeSnippetsService(/*set_empty_response=*/false);
 
   // When recreating the service, we should not get any |Schedule| calls:
   // The tasks are already scheduled with the correct intervals, so nothing on
   // initialization, and the service has data from the DB, so no automatic fetch
   // should happen.
   Mock::VerifyAndClearExpectations(&mock_scheduler());
   EXPECT_CALL(mock_scheduler(), Schedule(_, _)).Times(0);
   EXPECT_CALL(mock_scheduler(), Unschedule()).Times(0);
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
 }
 
 TEST_F(RemoteSuggestionsProviderTest, RescheduleAfterSuccessfulFetch) {
   // We should get two |Schedule| calls: The first when initialization
   // completes, the second one after the automatic (since the service doesn't
   // have any data yet) fetch finishes.
   EXPECT_CALL(mock_scheduler(), Schedule(_, _)).Times(2);
   auto service = MakeSnippetsService();
 
   // A successful fetch should trigger another |Schedule|.
@@ skipping 130 matching lines @@
   CategoryInfo info_with_title = service->GetCategoryInfo(articles_category());
   EXPECT_THAT(info_before.title(), Not(Eq(info_with_title.title())));
   EXPECT_THAT(test_default_title, Eq(info_with_title.title()));
   EXPECT_THAT(info_before.has_more_action(), Eq(true));
   EXPECT_THAT(info_before.has_reload_action(), Eq(true));
   EXPECT_THAT(info_before.has_view_all_action(), Eq(false));
   EXPECT_THAT(info_before.show_if_empty(), Eq(true));
 }
 
 TEST_F(RemoteSuggestionsProviderTest, MultipleCategories) {
-  std::string json_str(
-      GetMultiCategoryJson({GetSnippetN(0)}, {GetSnippetN(1)}));
-
   auto service = MakeSnippetsService();
-
+  std::string json_str =
+      MultiCategoryJsonBuilder()
+          .AddCategory({GetSnippetN(0)}, /*remote_category_id=*/1)
+          .AddCategory({GetSnippetN(1)}, /*remote_category_id=*/2)
+          .Build();
   LoadFromJSONString(service.get(), json_str);
 
   ASSERT_THAT(observer().statuses(),
               Eq(std::map<Category, CategoryStatus, Category::CompareByID>{
                   {articles_category(), CategoryStatus::AVAILABLE},
                   {other_category(), CategoryStatus::AVAILABLE},
               }));
 
   EXPECT_THAT(service->GetSnippetsForTesting(articles_category()), SizeIs(1));
   EXPECT_THAT(service->GetSnippetsForTesting(other_category()), SizeIs(1));
@@ skipping 32 matching lines @@
   auto service = MakeSnippetsService();
   CategoryInfo article_info = service->GetCategoryInfo(articles_category());
   EXPECT_THAT(article_info.has_more_action(), Eq(true));
   EXPECT_THAT(article_info.has_reload_action(), Eq(true));
   EXPECT_THAT(article_info.has_view_all_action(), Eq(false));
   EXPECT_THAT(article_info.show_if_empty(), Eq(true));
 }
 
 TEST_F(RemoteSuggestionsProviderTest, ExperimentalCategoryInfo) {
   auto service = MakeSnippetsService();
-
+  std::string json_str =
+      MultiCategoryJsonBuilder()
+          .AddCategory({GetSnippetN(0)}, /*remote_category_id=*/1)
+          .AddCategory({GetSnippetN(1)}, kUnknownRemoteCategoryId)
+          .Build();
   // Load data with multiple categories so that a new experimental category gets
   // registered.
-  LoadFromJSONString(service.get(),
-                     GetMultiCategoryJson({GetSnippetN(0)}, {GetSnippetN(1)},
-                                          kUnknownRemoteCategoryId));
+  LoadFromJSONString(service.get(), json_str);
+
   CategoryInfo info = service->GetCategoryInfo(unknown_category());
   EXPECT_THAT(info.has_more_action(), Eq(false));
   EXPECT_THAT(info.has_reload_action(), Eq(false));
   EXPECT_THAT(info.has_view_all_action(), Eq(false));
   EXPECT_THAT(info.show_if_empty(), Eq(false));
 }
 
+TEST_F(RemoteSuggestionsProviderTest, AddRemoteCategoriesToCategoryRanker) {
+  auto mock_ranker = base::MakeUnique<MockCategoryRanker>();
+  MockCategoryRanker* raw_mock_ranker = mock_ranker.get();
+  SetCategoryRanker(std::move(mock_ranker));
+  std::string json_str =
+      MultiCategoryJsonBuilder()
+          .AddCategory({GetSnippetN(0)}, /*remote_category_id=*/11)
+          .AddCategory({GetSnippetN(1)}, /*remote_category_id=*/13)
+          .AddCategory({GetSnippetN(2)}, /*remote_category_id=*/12)
+          .Build();
+  SetUpFetchResponse(json_str);
+  {
+    // The order of categories is determined by the order in which they are
+    // added. Thus, the latter is tested here.
+    InSequence s;
+    EXPECT_CALL(*raw_mock_ranker,
+                AppendCategoryIfNecessary(Category::FromRemoteCategory(11)));
+    EXPECT_CALL(*raw_mock_ranker,
+                AppendCategoryIfNecessary(Category::FromRemoteCategory(13)));
+    EXPECT_CALL(*raw_mock_ranker,
+                AppendCategoryIfNecessary(Category::FromRemoteCategory(12)));
+  }
+  auto service = MakeSnippetsService(/*set_empty_response=*/false);
+}
+
 TEST_F(RemoteSuggestionsProviderTest, PersistCategoryInfos) {
   auto service = MakeSnippetsService();
-
-  LoadFromJSONString(service.get(),
-                     GetMultiCategoryJson({GetSnippetN(0)}, {GetSnippetN(1)},
-                                          kUnknownRemoteCategoryId));
+  // TODO(vitaliii): Use |articles_category()| instead of constant ID below.
+  std::string json_str =
+      MultiCategoryJsonBuilder()
+          .AddCategoryWithCustomTitle(
+              {GetSnippetN(0)}, /*remote_category_id=*/1, "Articles for You")
+          .AddCategoryWithCustomTitle({GetSnippetN(1)},
+                                      kUnknownRemoteCategoryId, "Other Things")
+          .Build();
+  LoadFromJSONString(service.get(), json_str);
 
   ASSERT_EQ(observer().StatusForCategory(articles_category()),
             CategoryStatus::AVAILABLE);
   ASSERT_EQ(observer().StatusForCategory(unknown_category()),
             CategoryStatus::AVAILABLE);
 
   CategoryInfo info_articles_before =
       service->GetCategoryInfo(articles_category());
   CategoryInfo info_unknown_before =
       service->GetCategoryInfo(unknown_category());
 
   // Recreate the service to simulate a Chrome restart.
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
 
   // The categories should have been restored.
   ASSERT_NE(observer().StatusForCategory(articles_category()),
             CategoryStatus::NOT_PROVIDED);
   ASSERT_NE(observer().StatusForCategory(unknown_category()),
             CategoryStatus::NOT_PROVIDED);
 
   EXPECT_EQ(observer().StatusForCategory(articles_category()),
             CategoryStatus::AVAILABLE);
   EXPECT_EQ(observer().StatusForCategory(unknown_category()),
             CategoryStatus::AVAILABLE);
 
   CategoryInfo info_articles_after =
       service->GetCategoryInfo(articles_category());
   CategoryInfo info_unknown_after =
       service->GetCategoryInfo(unknown_category());
 
   EXPECT_EQ(info_articles_before.title(), info_articles_after.title());
   EXPECT_EQ(info_unknown_before.title(), info_unknown_after.title());
 }
 
+TEST_F(RemoteSuggestionsProviderTest, PersistRemoteCategoryOrder) {
+  // We create a service with a normal ranker to store the order.
+  std::string json_str =
+      MultiCategoryJsonBuilder()
+          .AddCategory({GetSnippetN(0)}, /*remote_category_id=*/11)
+          .AddCategory({GetSnippetN(1)}, /*remote_category_id=*/13)
+          .AddCategory({GetSnippetN(2)}, /*remote_category_id=*/12)
+          .Build();
+  SetUpFetchResponse(json_str);
+  auto service = MakeSnippetsService(/*set_empty_response=*/false);
+
+  // We manually recreate the service to simulate Chrome restart and enforce a
+  // mock ranker. The response is cleared to ensure that the order is not
+  // fetched.
+  SetUpFetchResponse("");
+  auto mock_ranker = base::MakeUnique<MockCategoryRanker>();
+  MockCategoryRanker* raw_mock_ranker = mock_ranker.get();
+  SetCategoryRanker(std::move(mock_ranker));
+  {
+    // The order of categories is determined by the order in which they are
+    // added. Thus, the latter is tested here.
+    InSequence s;
+    // Article category always exists and, therefore, it is stored in prefs too.
+    EXPECT_CALL(*raw_mock_ranker,
+                AppendCategoryIfNecessary(articles_category()));
+
+    EXPECT_CALL(*raw_mock_ranker,
+                AppendCategoryIfNecessary(Category::FromRemoteCategory(11)));
+    EXPECT_CALL(*raw_mock_ranker,
+                AppendCategoryIfNecessary(Category::FromRemoteCategory(13)));
+    EXPECT_CALL(*raw_mock_ranker,
+                AppendCategoryIfNecessary(Category::FromRemoteCategory(12)));
+  }
+  ResetSnippetsService(&service, /*set_empty_response=*/false);
+}
+
 TEST_F(RemoteSuggestionsProviderTest, PersistSuggestions) {
   auto service = MakeSnippetsService();
-
-  LoadFromJSONString(service.get(),
-                     GetMultiCategoryJson({GetSnippetN(0)}, {GetSnippetN(1)}));
+  std::string json_str =
+      MultiCategoryJsonBuilder()
+          .AddCategory({GetSnippetN(0)}, /*remote_category_id=*/1)
+          .AddCategory({GetSnippetN(2)}, /*remote_category_id=*/2)
+          .Build();
+  LoadFromJSONString(service.get(), json_str);
 
   ASSERT_THAT(observer().SuggestionsForCategory(articles_category()),
               SizeIs(1));
   ASSERT_THAT(observer().SuggestionsForCategory(other_category()), SizeIs(1));
 
   // Recreate the service to simulate a Chrome restart.
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
 
   // The suggestions in both categories should have been restored.
   EXPECT_THAT(observer().SuggestionsForCategory(articles_category()),
               SizeIs(1));
   EXPECT_THAT(observer().SuggestionsForCategory(other_category()), SizeIs(1));
 }
 
 TEST_F(RemoteSuggestionsProviderTest, DontNotifyIfNotAvailable) {
   // Get some suggestions into the database.
   auto service = MakeSnippetsService();
-  LoadFromJSONString(service.get(),
-                     GetMultiCategoryJson({GetSnippetN(0)}, {GetSnippetN(1)}));
+  std::string json_str =
+      MultiCategoryJsonBuilder()
+          .AddCategory({GetSnippetN(0)},
+                       /*remote_category_id=*/1)
+          .AddCategory({GetSnippetN(1)}, /*remote_category_id=*/2)
+          .Build();
+  LoadFromJSONString(service.get(), json_str);
+
   ASSERT_THAT(observer().SuggestionsForCategory(articles_category()),
               SizeIs(1));
   ASSERT_THAT(observer().SuggestionsForCategory(other_category()), SizeIs(1));
 
   service.reset();
 
   // Set the pref that disables remote suggestions.
   pref_service()->SetBoolean(prefs::kEnableSnippets, false);
 
   // Recreate the service to simulate a Chrome start.
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
 
   ASSERT_THAT(RemoteSuggestionsProvider::State::DISABLED, Eq(service->state_));
 
   // Now the observer should not have received any suggestions.
   EXPECT_THAT(observer().SuggestionsForCategory(articles_category()),
               IsEmpty());
   EXPECT_THAT(observer().SuggestionsForCategory(other_category()), IsEmpty());
 }
 
 TEST_F(RemoteSuggestionsProviderTest, Clear) {
@@ skipping 352 matching lines @@
   image_decoder()->SetDecodedImage(gfx::test::CreateImage(1, 1));
   image = FetchImage(service.get(), MakeArticleID(kSnippetUrl));
   EXPECT_FALSE(image.IsEmpty());
   EXPECT_EQ(1, image.Width());
 
   // Make sure that fetching the same snippet again does not re-add it.
   LoadFromJSONString(service.get(), json_str);
   EXPECT_THAT(service->GetSnippetsForTesting(articles_category()), IsEmpty());
 
   // The snippet should stay dismissed even after re-creating the service.
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
   LoadFromJSONString(service.get(), json_str);
   EXPECT_THAT(service->GetSnippetsForTesting(articles_category()), IsEmpty());
 
   // The snippet can be added again after clearing dismissed snippets.
   service->ClearDismissedSuggestionsForDebugging(articles_category());
   EXPECT_THAT(service->GetSnippetsForTesting(articles_category()), IsEmpty());
   LoadFromJSONString(service.get(), json_str);
   EXPECT_THAT(service->GetSnippetsForTesting(articles_category()), SizeIs(1));
 }
 
@@ skipping 175 matching lines @@
   EXPECT_THAT(tester.GetAllSamples("NewTabPage.Snippets.NumArticlesFetched"),
               ElementsAre(base::Bucket(/*min=*/0, /*count=*/1),
                           base::Bucket(/*min=*/1, /*count=*/3)));
   EXPECT_THAT(
       tester.GetAllSamples("NewTabPage.Snippets.NumArticlesZeroDueToDiscarded"),
       ElementsAre(base::Bucket(/*min=*/1, /*count=*/1)));
 
   // There is only a single, dismissed snippet in the database, so recreating
   // the service will require us to re-fetch.
   tester.ExpectTotalCount("NewTabPage.Snippets.NumArticlesFetched", 4);
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
   EXPECT_EQ(observer().StatusForCategory(articles_category()),
             CategoryStatus::AVAILABLE);
   tester.ExpectTotalCount("NewTabPage.Snippets.NumArticlesFetched", 5);
   EXPECT_THAT(
       tester.GetAllSamples("NewTabPage.Snippets.NumArticlesZeroDueToDiscarded"),
       ElementsAre(base::Bucket(/*min=*/1, /*count=*/2)));
 
   // But if there's a non-dismissed snippet in the database, recreating it
   // shouldn't trigger a fetch.
   LoadFromJSONString(
       service.get(),
       GetTestJson({GetSnippetWithUrl("http://not-dismissed.com")}));
   tester.ExpectTotalCount("NewTabPage.Snippets.NumArticlesFetched", 6);
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
   tester.ExpectTotalCount("NewTabPage.Snippets.NumArticlesFetched", 6);
 }
 
 TEST_F(RemoteSuggestionsProviderTest, DismissShouldRespectAllKnownUrls) {
   auto service = MakeSnippetsService();
 
   const base::Time creation = GetDefaultCreationTime();
   const base::Time expiry = GetDefaultExpirationTime();
   const std::vector<std::string> source_urls = {
       "http://mashable.com/2016/05/11/stolen",
@@ skipping 176 matching lines @@
   gfx::Image image = FetchImage(service.get(), MakeArticleID(kSnippetUrl));
   EXPECT_EQ(1, image.Width());
   EXPECT_FALSE(image.IsEmpty());
 
   // Send new suggestion which don't include the snippet referencing the image.
   LoadFromJSONString(service.get(),
                      GetTestJson({GetSnippetWithUrl(
                          "http://something.com/pletely/unrelated")}));
   // The image should still be available until a restart happens.
   EXPECT_FALSE(FetchImage(service.get(), MakeArticleID(kSnippetUrl)).IsEmpty());
-  ResetSnippetsService(&service);
+  ResetSnippetsService(&service, /*set_empty_response=*/true);
   // After the restart, the image should be garbage collected.
   EXPECT_TRUE(FetchImage(service.get(), MakeArticleID(kSnippetUrl)).IsEmpty());
 }
 
 TEST_F(RemoteSuggestionsProviderTest,
        ShouldHandleMoreThanMaxSnippetsInResponse) {
   auto service = MakeSnippetsService();
 
   std::vector<std::string> suggestions;
   for (int i = 0; i < service->GetMaxSnippetCountForTesting() + 1; ++i) {
@@ skipping 33 matching lines @@
   service->FetchSnippetsInTheBackground();
   base::RunLoop().RunUntilIdle();
   EXPECT_EQ(
       simple_test_clock_ptr->Now().ToInternalValue(),
       pref_service()->GetInt64(prefs::kLastSuccessfulBackgroundFetchTime));
   // TODO(markusheintz): Add a test that simulates a browser restart once the
   // scheduler refactoring is done (crbug.com/672434).
 }
 
 }  // namespace ntp_snippets