1 files changed, 57 insertions, 34 deletions

diff --git a/element_detection_unittest.cc b/element_detection_unittest.cc
index 2200c0b..6dbfa3f 100644
--- a/element_detection_unittest.cc
+++ b/element_detection_unittest.cc
@@ -11,14 +11,65 @@
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace zucchini {
-
 namespace {
+// This test uses a mock archive format where regions are determined by their
+// consecutive byte values rather than parsing real executables.
+//
+// 0 - Padding or raw data (not mapped to an executable).
+// 1 - A Win32x86 executable.
+// 2 - A Win32x64 executable.
+//
+// So an example archive file of;
+// 0 1 1 1 0 1 1 0 0 2 2 2 2
+// contains (in order left to right):
+// - One padding byte
+// - Three byte Win32x86 executable
+// - One padding byte
+// - Two byte Win32x86 executable
+// - Two padding bytes
+// - Four byte Win32x64 executable
 
-using ElementVector = std::vector<Element>;
+class ElementDetectionTest : public ::testing::Test {
+ protected:
+  using ElementVector = std::vector<Element>;
+  using ExeTypeMap = std::map<uint8_t, ExecutableType>;
 
-}  // namespace
+  ElementDetectionTest()
+      : exe_map_({{1, kExeTypeWin32X86}, {2, kExeTypeWin32X64}}) {}
+
+  ElementVector TestElementFinder(std::vector<uint8_t> buffer) {
+    ConstBufferView image(buffer.data(), buffer.size());
+
+    ElementFinder finder(
+        image,
+        base::BindRepeating(
+            [](ExeTypeMap exe_map, ConstBufferView image,
+               ConstBufferView region) -> base::Optional<Element> {
+              EXPECT_GE(region.begin(), image.begin());
+              EXPECT_LE(region.end(), image.end());
+              EXPECT_GE(region.size(), 0U);
+
+              if (region[0] != 0) {
+                offset_t length = 1;
+                while (length < region.size() && region[length] == region[0])
+                  ++length;
+                return Element{{0, length}, exe_map[region[0]]};
+              }
+              return base::nullopt;
+            },
+            exe_map_, image));
+    std::vector<Element> elements;
+    for (auto element = finder.GetNext(); element; element = finder.GetNext()) {
+      elements.push_back(*element);
+    }
+    return elements;
+  }
+
+  // Translation map from mock archive bytes to actual types used in Zucchini.
+  ExeTypeMap exe_map_;
+};
 
-TEST(ElementDetectionTest, ElementFinderEmpty) {
+TEST_F(ElementDetectionTest, ElementFinderEmpty) {
   std::vector<uint8_t> buffer(10, 0);
   ElementFinder finder(
       ConstBufferView(buffer.data(), buffer.size()),
@@ -28,36 +79,7 @@ TEST(ElementDetectionTest, ElementFinderEmpty) {
   EXPECT_EQ(base::nullopt, finder.GetNext());
 }
 
-ElementVector TestElementFinder(std::vector<uint8_t> buffer) {
-  ConstBufferView image(buffer.data(), buffer.size());
-
-  ElementFinder finder(
-      image,
-      base::BindRepeating(
-          [](ConstBufferView image,
-             ConstBufferView region) -> base::Optional<Element> {
-            EXPECT_GE(region.begin(), image.begin());
-            EXPECT_LE(region.end(), image.end());
-            EXPECT_GE(region.size(), 0U);
-
-            if (region[0] != 0) {
-              offset_t length = 1;
-              while (length < region.size() && region[length] == region[0])
-                ++length;
-              return Element{{0, length},
-                             static_cast<ExecutableType>(region[0])};
-            }
-            return base::nullopt;
-          },
-          image));
-  std::vector<Element> elements;
-  for (auto element = finder.GetNext(); element; element = finder.GetNext()) {
-    elements.push_back(*element);
-  }
-  return elements;
-}
-
-TEST(ElementDetectionTest, ElementFinder) {
+TEST_F(ElementDetectionTest, ElementFinder) {
   EXPECT_EQ(ElementVector(), TestElementFinder({}));
   EXPECT_EQ(ElementVector(), TestElementFinder({0, 0}));
   EXPECT_EQ(ElementVector({{{0, 2}, kExeTypeWin32X86}}),
@@ -75,4 +97,5 @@ TEST(ElementDetectionTest, ElementFinder) {
       TestElementFinder({0, 1, 1, 0, 2, 2, 2}));
 }
 
+}  // namespace
 }  // namespace zucchini