1 files changed, 97 insertions, 0 deletions
diff --git a/toolkit/components/url-classifier/tests/gtest/TestFailUpdate.cpp b/toolkit/components/url-classifier/tests/gtest/TestFailUpdate.cpp
new file mode 100644
index 000000000..bdb9eebb0
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/gtest/TestFailUpdate.cpp
@@ -0,0 +1,97 @@
+#include "HashStore.h"
+#include "nsPrintfCString.h"
+#include "string.h"
+#include "gtest/gtest.h"
+#include "mozilla/Unused.h"
+
+using namespace mozilla;
+using namespace mozilla::safebrowsing;
+
+static const char* kFilesInV2[] = {".pset", ".sbstore"};
+static const char* kFilesInV4[] = {".pset", ".metadata"};
+
+#define V2_TABLE  "gtest-malware-simple"
+#define V4_TABLE1 "goog-malware-proto"
+#define V4_TABLE2 "goog-phish-proto"
+
+#define ROOT_DIR        NS_LITERAL_STRING("safebrowsing")
+#define SB_FILE(x, y)   NS_ConvertUTF8toUTF16(nsPrintfCString("%s%s",x, y))
+
+template<typename T, size_t N>
+void CheckFileExist(const char* table, const T (&files)[N], bool expectExists)
+{
+  for (uint32_t i = 0; i < N; i++) {
+    // This is just a quick way to know if this is v4 table
+    NS_ConvertUTF8toUTF16 SUB_DIR(strstr(table, "-proto") ? "google4" : "");
+    nsCOMPtr<nsIFile> file =
+      GetFile(nsTArray<nsString> { ROOT_DIR, SUB_DIR, SB_FILE(table, files[i]) });
+
+    bool exists;
+    file->Exists(&exists);
+
+    nsAutoCString path;
+    file->GetNativePath(path);
+    ASSERT_EQ(expectExists, exists) << path.get();
+  }
+}
+
+TEST(FailUpdate, CheckTableReset)
+{
+  const bool FULL_UPDATE = true;
+  const bool PARTIAL_UPDATE = false;
+
+  // Apply V2 update
+  {
+    auto update = new TableUpdateV2(NS_LITERAL_CSTRING(V2_TABLE));
+    Unused << update->NewAddChunk(1);
+
+    ApplyUpdate(update);
+
+    // A successful V2 update should create .pset & .sbstore files
+    CheckFileExist(V2_TABLE, kFilesInV2, true);
+  }
+
+  // Helper function to generate table update data
+  auto func = [](TableUpdateV4* update, bool full, const char* str) {
+    update->SetFullUpdate(full);
+    std::string prefix(str);
+    update->NewPrefixes(prefix.length(), prefix);
+  };
+
+  // Apply V4 update for table1
+  {
+    auto update = new TableUpdateV4(NS_LITERAL_CSTRING(V4_TABLE1));
+    func(update, FULL_UPDATE, "test_prefix");
+
+    ApplyUpdate(update);
+
+    // A successful V4 update should create .pset & .metadata files
+    CheckFileExist(V4_TABLE1, kFilesInV4, true);
+  }
+
+  // Apply V4 update for table2
+  {
+    auto update = new TableUpdateV4(NS_LITERAL_CSTRING(V4_TABLE2));
+    func(update, FULL_UPDATE, "test_prefix");
+
+    ApplyUpdate(update);
+
+    CheckFileExist(V4_TABLE2, kFilesInV4, true);
+  }
+
+  // Apply V4 update with the same prefix in previous full udpate
+  // This should cause an update error.
+  {
+    auto update = new TableUpdateV4(NS_LITERAL_CSTRING(V4_TABLE1));
+    func(update, PARTIAL_UPDATE, "test_prefix");
+
+    ApplyUpdate(update);
+
+    // A fail update should remove files for that table
+    CheckFileExist(V4_TABLE1, kFilesInV4, false);
+
+    // A fail update should NOT remove files for the other tables
+    CheckFileExist(V2_TABLE, kFilesInV2, true);
+    CheckFileExist(V4_TABLE2, kFilesInV4, true);
+  }
+}