1 files changed, 87 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/gtest/TestPrefixSet.cpp b/toolkit/components/url-classifier/tests/gtest/TestPrefixSet.cpp
new file mode 100644
index 0000000000..6d83cd02f6
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/gtest/TestPrefixSet.cpp
@@ -0,0 +1,87 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "mozilla/Preferences.h"
+#include "nsString.h"
+#include "nsUrlClassifierPrefixSet.h"
+
+#include "Common.h"
+
+// This function generate N 4 byte prefixes.
+static void RandomPrefixes(uint32_t N, nsTArray<uint32_t>& array) {
+  array.Clear();
+  array.SetCapacity(N);
+
+  for (uint32_t i = 0; i < N; i++) {
+    bool added = false;
+
+    while (!added) {
+      nsAutoCString prefix;
+      char* dst = prefix.BeginWriting();
+      for (uint32_t j = 0; j < 4; j++) {
+        dst[j] = static_cast<char>(rand() % 256);
+      }
+
+      const char* src = prefix.BeginReading();
+      uint32_t data = 0;
+      memcpy(&data, src, sizeof(data));
+
+      if (!array.Contains(data)) {
+        array.AppendElement(data);
+        added = true;
+      }
+    }
+  }
+
+  struct Comparator {
+    bool LessThan(const uint32_t& aA, const uint32_t& aB) const {
+      return aA < aB;
+    }
+
+    bool Equals(const uint32_t& aA, const uint32_t& aB) const {
+      return aA == aB;
+    }
+  };
+
+  array.Sort(Comparator());
+}
+
+void RunTest(uint32_t aTestSize) {
+  RefPtr<nsUrlClassifierPrefixSet> prefixSet = new nsUrlClassifierPrefixSet();
+  nsTArray<uint32_t> array;
+
+  RandomPrefixes(aTestSize, array);
+
+  nsresult rv = prefixSet->SetPrefixes(array.Elements(), array.Length());
+  ASSERT_NS_SUCCEEDED(rv);
+
+  for (uint32_t i = 0; i < array.Length(); i++) {
+    uint32_t value = 0;
+    rv = prefixSet->GetPrefixByIndex(i, &value);
+    ASSERT_NS_SUCCEEDED(rv);
+
+    ASSERT_TRUE(value == array[i]);
+  }
+}
+
+TEST(URLClassifierPrefixSet, GetTargetPrefixWithLargeSet)
+{
+  // Make sure the delta algorithm will be used.
+  static const char prefKey[] = "browser.safebrowsing.prefixset_max_array_size";
+  mozilla::Preferences::SetUint(prefKey, 10000);
+
+  // Ideally, we should test more than 512 * 1024 entries. But, it will make the
+  // test too long. So, we test 100k entries instead.
+  RunTest(100000);
+}
+
+TEST(URLClassifierPrefixSet, GetTargetPrefixWithSmallSet)
+{
+  // Make sure the delta algorithm won't be used.
+  static const char prefKey[] = "browser.safebrowsing.prefixset_max_array_size";
+  mozilla::Preferences::SetUint(prefKey, 10000);
+
+  RunTest(1000);
+}