Merge the hash prefixes from the old store and the additions in the partial

components/safe_browsing_db/v4_store_unittest.cc

Index: components/safe_browsing_db/v4_store_unittest.cc
diff --git a/components/safe_browsing_db/v4_store_unittest.cc b/components/safe_browsing_db/v4_store_unittest.cc
index ec41208a73e923cc979744df2f9087d220df38e8..353a848084ac1819322ad158df892cf8e6fb579c 100644
--- a/components/safe_browsing_db/v4_store_unittest.cc
+++ b/components/safe_browsing_db/v4_store_unittest.cc
@@ -141,4 +141,107 @@ TEST_F(V4StoreTest, TestWriteFullResponseType) {
   EXPECT_EQ("test_client_state", read_store->state_);
 }
 
+TEST_F(V4StoreTest, TestAddUnlumpedHashesWithInvalidAddition) {
+  HashPrefixMap prefix_map;
+  MergeUpdateResult result = V4Store::AddUnlumpedHashes(5, "a", &prefix_map);
+  EXPECT_EQ(ADDITIONS_SIZE_UNEXPECTED_FAILURE, result);
+  EXPECT_TRUE(prefix_map.empty());
+}
+
+TEST_F(V4StoreTest, TestAddUnlumpedHashesWithEmptyString) {
+  HashPrefixMap prefix_map;
+  MergeUpdateResult result = V4Store::AddUnlumpedHashes(5, "", &prefix_map);
+  EXPECT_EQ(MERGE_SUCCESS, result);
+  EXPECT_TRUE(prefix_map.empty());
+}
+
+TEST_F(V4StoreTest, TestAddUnlumpedHashes) {
+  HashPrefixMap prefix_map;
+  EXPECT_EQ(MERGE_SUCCESS,
+            V4Store::AddUnlumpedHashes(5, "abcde5432100000-----", &prefix_map));
+  EXPECT_EQ(1u, prefix_map.size());
+  HashPrefixes& hash_prefixes = prefix_map[5];
+  EXPECT_EQ(4u, hash_prefixes.size());
+
+  EXPECT_EQ("abcde", std::string(hash_prefixes[0].get()));
+  EXPECT_EQ("54321", std::string(hash_prefixes[1].get()));
+  EXPECT_EQ("00000", std::string(hash_prefixes[2].get()));
+  EXPECT_EQ("-----", std::string(hash_prefixes[3].get()));
+
+  EXPECT_EQ(MERGE_SUCCESS,
+            V4Store::AddUnlumpedHashes(4, "abcde5432100000-----", &prefix_map));
+  EXPECT_EQ(2u, prefix_map.size());
+  HashPrefixes& hash_prefixes_of_length_4 = prefix_map[4];
+  EXPECT_EQ(5u, hash_prefixes_of_length_4.size());
+  EXPECT_EQ("abcd", std::string(hash_prefixes_of_length_4[0].get()));
+  EXPECT_EQ("e543", std::string(hash_prefixes_of_length_4[1].get()));
+  EXPECT_EQ("2100", std::string(hash_prefixes_of_length_4[2].get()));
+  EXPECT_EQ("000-", std::string(hash_prefixes_of_length_4[3].get()));
+  EXPECT_EQ("----", std::string(hash_prefixes_of_length_4[4].get()));
+}
+
+TEST_F(V4StoreTest, TestGetNextSmallestPrefixSizeWithEmptyPrefixMap) {
+  HashPrefixMap prefix_map;
+  CounterMap counter_map = V4Store::GetInitializedCounterMap(prefix_map);
+
+  PrefixSize prefix_size;
+  EXPECT_FALSE(V4Store::GetNextSmallestPrefixSize(prefix_map, counter_map,
+                                                  &prefix_size));
+}
+
+TEST_F(V4StoreTest, TestGetNextSmallestPrefixSize) {
+  HashPrefixMap prefix_map;
+  V4Store::AddUnlumpedHashes(5, "abcde5432100000-----", &prefix_map);
+  V4Store::AddUnlumpedHashes(4, "abcde5432100000-----", &prefix_map);
+  CounterMap counter_map = V4Store::GetInitializedCounterMap(prefix_map);
+
+  PrefixSize prefix_size;
+  EXPECT_TRUE(V4Store::GetNextSmallestPrefixSize(prefix_map, counter_map,
+                                                 &prefix_size));
+  EXPECT_EQ(4u, prefix_size);
+}
+
+TEST_F(V4StoreTest, TestGetNextUnmergedPrefix) {
+  HashPrefixMap prefix_map;
+  V4Store::AddUnlumpedHashes(5, "abcde5432100000-----", &prefix_map);
+  V4Store::AddUnlumpedHashes(4, "abcde5432100000-----", &prefix_map);
+  CounterMap counter_map = V4Store::GetInitializedCounterMap(prefix_map);
+
+  PrefixSize prefix_size;
+  EXPECT_TRUE(V4Store::GetNextSmallestPrefixSize(prefix_map, counter_map,
+                                                 &prefix_size));
+  const HashPrefix& prefix = V4Store::GetNextUnmergedPrefixForSize(
+      prefix_size, prefix_map, counter_map);
+  std::string prefix_str = std::string(prefix.get(), prefix_size);
+  EXPECT_EQ("abcd", prefix_str);
+}
+
+TEST_F(V4StoreTest, TestMergeUpdates) {
+  HashPrefixMap prefix_map_old;
+  V4Store::AddUnlumpedHashes(4, "abcdefgh", &prefix_map_old);
+  V4Store::AddUnlumpedHashes(5, "54321abcde", &prefix_map_old);
+  HashPrefixMap prefix_map_additions;
+  V4Store::AddUnlumpedHashes(4, "----1111bbbb", &prefix_map_additions);
+  V4Store::AddUnlumpedHashes(5, "22222bcdef", &prefix_map_additions);
+
+  std::unique_ptr<V4Store> store(new V4Store(task_runner_, store_path_));
+  store->MergeUpdate(prefix_map_old, prefix_map_additions);
+  const HashPrefixMap& prefix_map = store->hash_prefix_map_;
+  EXPECT_EQ(2u, prefix_map.size());
+  const HashPrefixes& hash_prefixes_of_length_4 = prefix_map.at(4);
+  EXPECT_EQ(5u, hash_prefixes_of_length_4.size());
+  EXPECT_EQ("----", std::string(hash_prefixes_of_length_4[0].get()));
+  EXPECT_EQ("1111", std::string(hash_prefixes_of_length_4[1].get()));
+  EXPECT_EQ("abcd", std::string(hash_prefixes_of_length_4[2].get()));
+  EXPECT_EQ("bbbb", std::string(hash_prefixes_of_length_4[3].get()));
+  EXPECT_EQ("efgh", std::string(hash_prefixes_of_length_4[4].get()));
+
+  const HashPrefixes& hash_prefixes_of_length_5 = prefix_map.at(5);
+  EXPECT_EQ(4u, hash_prefixes_of_length_5.size());
+  EXPECT_EQ("22222", std::string(hash_prefixes_of_length_5[0].get()));
+  EXPECT_EQ("54321", std::string(hash_prefixes_of_length_5[1].get()));
+  EXPECT_EQ("abcde", std::string(hash_prefixes_of_length_5[2].get()));
+  EXPECT_EQ("bcdef", std::string(hash_prefixes_of_length_5[3].get()));
+}
+
 }  // namespace safe_browsing