CodeGen: refactor unit tests using parameterized tests

sandbox/linux/seccomp-bpf/codegen_unittest.cc

Index: sandbox/linux/seccomp-bpf/codegen_unittest.cc
diff --git a/sandbox/linux/seccomp-bpf/codegen_unittest.cc b/sandbox/linux/seccomp-bpf/codegen_unittest.cc
index d4760a56077ca5e4dcdb4d2f1fdeedf729dc784c..d7cbecf7720b4e50700ee3791d00cce58512839f 100644
--- a/sandbox/linux/seccomp-bpf/codegen_unittest.cc
+++ b/sandbox/linux/seccomp-bpf/codegen_unittest.cc
@@ -4,7 +4,6 @@
 
 #include "sandbox/linux/seccomp-bpf/codegen.h"
 
-#include <errno.h>
 #include <linux/filter.h>
 
 #include <set>
@@ -14,8 +13,7 @@
 #include "sandbox/linux/seccomp-bpf/basicblock.h"
 #include "sandbox/linux/seccomp-bpf/errorcode.h"
 #include "sandbox/linux/seccomp-bpf/instruction.h"
-#include "sandbox/linux/seccomp-bpf/sandbox_bpf.h"
-#include "sandbox/linux/tests/unit_tests.h"
+#include "testing/gtest/include/gtest/gtest.h"
 
 namespace sandbox {
 
@@ -23,58 +21,101 @@ namespace sandbox {
 // do so by defining a "friend" that makes these methods public for us.
 class CodeGenUnittestHelper : public CodeGen {
  public:
-  void FindBranchTargets(const Instruction& instructions,
-                         BranchTargets* branch_targets) {
-    CodeGen::FindBranchTargets(instructions, branch_targets);
+  using CodeGen::CutGraphIntoBasicBlocks;
+  using CodeGen::FindBranchTargets;
+  using CodeGen::MergeTails;
+};
+
+namespace {
+
+enum {
+  NO_FLAGS = 0x0000,
+  HAS_MERGEABLE_TAILS = 0x0001,
+};
+
+using ProgramTestFunc = void (*)(CodeGenUnittestHelper* gen,
+                                 Instruction* head,
+                                 int flags);
+
+class ProgramTest : public ::testing::TestWithParam<ProgramTestFunc> {
+ protected:
+  ProgramTest() : gen_() {}
+
+  // RunTest runs the test function argument.  It should be called at
+  // the end of each program test case.
+  void RunTest(Instruction* head, int flags) { GetParam()(&gen_, head, flags); }
+
+  Instruction* MakeInstruction(uint16_t code,
+                               uint32_t k,
+                               Instruction* next = nullptr) {
+    Instruction* ret = gen_.MakeInstruction(code, k, next);
+    EXPECT_TRUE(ret);

[jln (very slow on Chromium), 2014/11/03 20:19:56]

EXPECT_NE?

Also this should be an ASSERT....

...but, one issue though is that with GTest, ASSERTs only leave the current
function, but the calling function will continue to run normally (that's because
they can't use exceptions).

With SANDBOX_TEST, we worked around this since our assert could simply _exit().

[mdempsky, 2014/11/03 20:43:00]

Done.
Yeah, the other issue is ASSERT can't be used in a function that returns
non-void, hence why I'm using EXPECT here. :(

+    EXPECT_EQ(code, ret->code);
+    EXPECT_EQ(k, ret->k);
+    if (code == BPF_JMP + BPF_JA) {
+      // Annoying inconsistency.
+      EXPECT_EQ(nullptr, ret->next);
+      EXPECT_EQ(next, ret->jt_ptr);
+    } else {
+      EXPECT_EQ(next, ret->next);
+      EXPECT_EQ(nullptr, ret->jt_ptr);
+    }
+    EXPECT_EQ(nullptr, ret->jf_ptr);
+    return ret;
   }
 
-  BasicBlock* CutGraphIntoBasicBlocks(Instruction* insns,
-                                      const BranchTargets& branch_targets,
-                                      TargetsToBlocks* blocks) {
-    return CodeGen::CutGraphIntoBasicBlocks(insns, branch_targets, blocks);
+  Instruction* MakeInstruction(uint16_t code,
+                               uint32_t k,
+                               Instruction* jt,
+                               Instruction* jf) {
+    Instruction* ret = gen_.MakeInstruction(code, k, jt, jf);
+    EXPECT_TRUE(ret);

[jln (very slow on Chromium), 2014/11/03 20:19:56]

EXPECT_NE?

[mdempsky, 2014/11/03 20:43:00]

Done.

+    EXPECT_EQ(code, ret->code);
+    EXPECT_EQ(k, ret->k);
+    EXPECT_EQ(nullptr, ret->next);
+    EXPECT_EQ(jt, ret->jt_ptr);
+    EXPECT_EQ(jf, ret->jf_ptr);
+    return ret;
   }
 
-  void MergeTails(TargetsToBlocks* blocks) { CodeGen::MergeTails(blocks); }
+ private:
+  CodeGenUnittestHelper gen_;
 };
 
-enum { NO_FLAGS = 0x0000, HAS_MERGEABLE_TAILS = 0x0001, };
-
-Instruction* SampleProgramOneInstruction(CodeGen* codegen, int* flags) {
+TEST_P(ProgramTest, OneInstruction) {
   // Create the most basic valid BPF program:
   //    RET 0
-  *flags = NO_FLAGS;
-  return codegen->MakeInstruction(BPF_RET + BPF_K, 0);
+  Instruction* head = MakeInstruction(BPF_RET + BPF_K, 0);
+  RunTest(head, NO_FLAGS);
 }
 
-Instruction* SampleProgramSimpleBranch(CodeGen* codegen, int* flags) {
+TEST_P(ProgramTest, SimpleBranch) {
   // Create a program with a single branch:
   //    JUMP if eq 42 then $0 else $1
   // 0: RET 1
   // 1: RET 0
-  *flags = NO_FLAGS;
-  return codegen->MakeInstruction(
-      BPF_JMP + BPF_JEQ + BPF_K,
-      42,
-      codegen->MakeInstruction(BPF_RET + BPF_K, 1),
-      codegen->MakeInstruction(BPF_RET + BPF_K, 0));
+  Instruction* head = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K,
+                                      42,
+                                      MakeInstruction(BPF_RET + BPF_K, 1),
+                                      MakeInstruction(BPF_RET + BPF_K, 0));
+  RunTest(head, NO_FLAGS);
 }
 
-Instruction* SampleProgramAtypicalBranch(CodeGen* codegen, int* flags) {
+TEST_P(ProgramTest, AtypicalBranch) {
   // Create a program with a single branch:
   //    JUMP if eq 42 then $0 else $0
   // 0: RET 0
 
+  Instruction* ret = MakeInstruction(BPF_RET + BPF_K, 0);
+  Instruction* head = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 42, ret, ret);
+
   // N.B.: As the instructions in both sides of the branch are already
   //       the same object, we do not actually have any "mergeable" branches.
   //       This needs to be reflected in our choice of "flags".
-  *flags = NO_FLAGS;
-
-  Instruction* ret = codegen->MakeInstruction(
-      BPF_RET + BPF_K, 0);
-  return codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 42, ret, ret);
+  RunTest(head, NO_FLAGS);
 }
 
-Instruction* SampleProgramComplex(CodeGen* codegen, int* flags) {
+TEST_P(ProgramTest, Complex) {
   // Creates a basic BPF program that we'll use to test some of the code:
   //    JUMP if eq 42 the $0 else $1     (insn6)
   // 0: LD 23                            (insn5)
@@ -84,46 +125,18 @@ Instruction* SampleProgramComplex(CodeGen* codegen, int* flags) {
   //    RET 42                           (insn0)
   // 4: LD 42                            (insn3)
   //    RET 42                           (insn3+)
-  *flags = HAS_MERGEABLE_TAILS;
-
-  Instruction* insn0 = codegen->MakeInstruction(BPF_RET + BPF_K, 42);
-  SANDBOX_ASSERT(insn0);
-  SANDBOX_ASSERT(insn0->code == BPF_RET + BPF_K);
-  SANDBOX_ASSERT(insn0->next == NULL);
-
-  Instruction* insn1 =
-      codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 42, insn0);
-  SANDBOX_ASSERT(insn1);
-  SANDBOX_ASSERT(insn1->code == BPF_LD + BPF_W + BPF_ABS);
-  SANDBOX_ASSERT(insn1->k == 42);
-  SANDBOX_ASSERT(insn1->next == insn0);
-
-  Instruction* insn2 = codegen->MakeInstruction(BPF_JMP + BPF_JA, 0, insn1);
-  SANDBOX_ASSERT(insn2);
-  SANDBOX_ASSERT(insn2->code == BPF_JMP + BPF_JA);
-  SANDBOX_ASSERT(insn2->jt_ptr == insn1);
+  Instruction* insn0 = MakeInstruction(BPF_RET + BPF_K, 42);
+  Instruction* insn1 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 42, insn0);
+  Instruction* insn2 = MakeInstruction(BPF_JMP + BPF_JA, 0, insn1);
 
   // We explicitly duplicate instructions so that MergeTails() can coalesce
   // them later.
-  Instruction* insn3 = codegen->MakeInstruction(
-      BPF_LD + BPF_W + BPF_ABS,
-      42,
-      codegen->MakeInstruction(BPF_RET + BPF_K, 42));
+  Instruction* insn3 = MakeInstruction(
+      BPF_LD + BPF_W + BPF_ABS, 42, MakeInstruction(BPF_RET + BPF_K, 42));
 
   Instruction* insn4 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 42, insn2, insn3);
-  SANDBOX_ASSERT(insn4);
-  SANDBOX_ASSERT(insn4->code == BPF_JMP + BPF_JEQ + BPF_K);
-  SANDBOX_ASSERT(insn4->k == 42);
-  SANDBOX_ASSERT(insn4->jt_ptr == insn2);
-  SANDBOX_ASSERT(insn4->jf_ptr == insn3);
-
-  Instruction* insn5 =
-      codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 23, insn4);
-  SANDBOX_ASSERT(insn5);
-  SANDBOX_ASSERT(insn5->code == BPF_LD + BPF_W + BPF_ABS);
-  SANDBOX_ASSERT(insn5->k == 23);
-  SANDBOX_ASSERT(insn5->next == insn4);
+      MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 42, insn2, insn3);
+  Instruction* insn5 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 23, insn4);
 
   // Force a basic block that ends in neither a jump instruction nor a return
   // instruction. It only contains "insn5". This exercises one of the less
@@ -132,12 +145,12 @@ Instruction* SampleProgramComplex(CodeGen* codegen, int* flags) {
   // another aspect of the topo-sort algorithm (namely, the ability to
   // correctly count the incoming branches for subtrees that are not disjunct).
   Instruction* insn6 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 42, insn5, insn4);
+      MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 42, insn5, insn4);
 
-  return insn6;
+  RunTest(insn6, HAS_MERGEABLE_TAILS);
 }
 
-Instruction* SampleProgramConfusingTails(CodeGen* codegen, int* flags) {
+TEST_P(ProgramTest, ConfusingTails) {
   // This simple program demonstrates https://crbug.com/351103/
   // The two "LOAD 0" instructions are blocks of their own. MergeTails() could
   // be tempted to merge them since they are the same. However, they are
@@ -155,24 +168,20 @@ Instruction* SampleProgramConfusingTails(CodeGen* codegen, int* flags) {
   //  5) if A == 0x1; then JMP 6 else JMP 7
   //  6) RET 0
   //  7) RET 1
-  *flags = NO_FLAGS;
-
-  Instruction* i7 = codegen->MakeInstruction(BPF_RET + BPF_K, 1);
-  Instruction* i6 = codegen->MakeInstruction(BPF_RET + BPF_K, 0);
-  Instruction* i5 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i6, i7);
-  Instruction* i4 = codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i5);
-  Instruction* i3 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 2, i4, i5);
-  Instruction* i2 = codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i3);
-  Instruction* i1 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i2, i3);
-  Instruction* i0 = codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 1, i1);
-
-  return i0;
+
+  Instruction* i7 = MakeInstruction(BPF_RET + BPF_K, 1);
+  Instruction* i6 = MakeInstruction(BPF_RET + BPF_K, 0);
+  Instruction* i5 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i6, i7);
+  Instruction* i4 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i5);
+  Instruction* i3 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 2, i4, i5);
+  Instruction* i2 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i3);
+  Instruction* i1 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i2, i3);
+  Instruction* i0 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 1, i1);
+
+  RunTest(i0, NO_FLAGS);
 }
 
-Instruction* SampleProgramConfusingTailsBasic(CodeGen* codegen, int* flags) {
+TEST_P(ProgramTest, ConfusingTailsBasic) {
   // Without the fix for https://crbug.com/351103/, (see
   // SampleProgramConfusingTails()), this would generate a cyclic graph and
   // crash as the two "LOAD 0" instructions would get merged.
@@ -183,22 +192,18 @@ Instruction* SampleProgramConfusingTailsBasic(CodeGen* codegen, int* flags) {
   // 3) if A == 0x2; then JMP 4 else JMP 5
   // 4) LOAD 0  // System call number
   // 5) RET 1
-  *flags = NO_FLAGS;
-
-  Instruction* i5 = codegen->MakeInstruction(BPF_RET + BPF_K, 1);
-  Instruction* i4 = codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i5);
-  Instruction* i3 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 2, i4, i5);
-  Instruction* i2 = codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i3);
-  Instruction* i1 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i2, i3);
-  Instruction* i0 = codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 1, i1);
-
-  return i0;
+
+  Instruction* i5 = MakeInstruction(BPF_RET + BPF_K, 1);
+  Instruction* i4 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i5);
+  Instruction* i3 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 2, i4, i5);
+  Instruction* i2 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 0, i3);
+  Instruction* i1 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i2, i3);
+  Instruction* i0 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 1, i1);
+
+  RunTest(i0, NO_FLAGS);
 }
 
-Instruction* SampleProgramConfusingTailsMergeable(CodeGen* codegen,
-                                                  int* flags) {
+TEST_P(ProgramTest, ConfusingTailsMergeable) {
   // This is similar to SampleProgramConfusingTails(), except that
   // instructions 2 and 4 are now RET instructions.
   // In PointerCompare(), this exercises the path where two blocks are of the
@@ -213,39 +218,17 @@ Instruction* SampleProgramConfusingTailsMergeable(CodeGen* codegen,
   // 5) if A == 0x1; then JMP 6 else JMP 7
   // 6) RET 0
   // 7) RET 1
-  *flags = HAS_MERGEABLE_TAILS;
-
-  Instruction* i7 = codegen->MakeInstruction(BPF_RET + BPF_K, 1);
-  Instruction* i6 = codegen->MakeInstruction(BPF_RET + BPF_K, 0);
-  Instruction* i5 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i6, i7);
-  Instruction* i4 = codegen->MakeInstruction(BPF_RET + BPF_K, 42);
-  Instruction* i3 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 2, i4, i5);
-  Instruction* i2 = codegen->MakeInstruction(BPF_RET + BPF_K, 42);
-  Instruction* i1 =
-      codegen->MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i2, i3);
-  Instruction* i0 = codegen->MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 1, i1);
-
-  return i0;
-}
-void ForAllPrograms(void (*test)(CodeGenUnittestHelper*, Instruction*, int)) {
-  Instruction* (*function_table[])(CodeGen* codegen, int* flags) = {
-    SampleProgramOneInstruction,
-    SampleProgramSimpleBranch,
-    SampleProgramAtypicalBranch,
-    SampleProgramComplex,
-    SampleProgramConfusingTails,
-    SampleProgramConfusingTailsBasic,
-    SampleProgramConfusingTailsMergeable,
-  };
-
-  for (size_t i = 0; i < arraysize(function_table); ++i) {
-    CodeGenUnittestHelper codegen;
-    int flags = NO_FLAGS;
-    Instruction *prg = function_table[i](&codegen, &flags);
-    test(&codegen, prg, flags);
-  }
+
+  Instruction* i7 = MakeInstruction(BPF_RET + BPF_K, 1);
+  Instruction* i6 = MakeInstruction(BPF_RET + BPF_K, 0);
+  Instruction* i5 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i6, i7);
+  Instruction* i4 = MakeInstruction(BPF_RET + BPF_K, 42);
+  Instruction* i3 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 2, i4, i5);
+  Instruction* i2 = MakeInstruction(BPF_RET + BPF_K, 42);
+  Instruction* i1 = MakeInstruction(BPF_JMP + BPF_JEQ + BPF_K, 1, i2, i3);
+  Instruction* i0 = MakeInstruction(BPF_LD + BPF_W + BPF_ABS, 1, i1);
+
+  RunTest(i0, HAS_MERGEABLE_TAILS);
 }
 
 void MakeInstruction(CodeGenUnittestHelper* codegen,
@@ -253,10 +236,6 @@ void MakeInstruction(CodeGenUnittestHelper* codegen,
   // Nothing to do here
 }
 
-SANDBOX_TEST(CodeGen, MakeInstruction) {
-  ForAllPrograms(MakeInstruction);
-}
-
 void FindBranchTargets(CodeGenUnittestHelper* codegen, Instruction* prg, int) {
   BranchTargets branch_targets;
   codegen->FindBranchTargets(*prg, &branch_targets);
@@ -275,28 +254,28 @@ void FindBranchTargets(CodeGenUnittestHelper* codegen, Instruction* prg, int) {
     all_instructions.insert(insn);
     if (BPF_CLASS(insn->code) == BPF_JMP) {
       target_instructions.insert(insn->jt_ptr);
-      SANDBOX_ASSERT(insn->jt_ptr != NULL);
-      SANDBOX_ASSERT(branch_targets.find(insn->jt_ptr) != end);
+      ASSERT_TRUE(insn->jt_ptr != NULL);
+      ASSERT_TRUE(branch_targets.find(insn->jt_ptr) != end);
       if (BPF_OP(insn->code) != BPF_JA) {
         target_instructions.insert(insn->jf_ptr);
-        SANDBOX_ASSERT(insn->jf_ptr != NULL);
-        SANDBOX_ASSERT(branch_targets.find(insn->jf_ptr) != end);
+        ASSERT_TRUE(insn->jf_ptr != NULL);
+        ASSERT_TRUE(branch_targets.find(insn->jf_ptr) != end);
         stack.push_back(insn->jf_ptr);
       }
       insn = insn->jt_ptr;
     } else if (BPF_CLASS(insn->code) == BPF_RET) {
-      SANDBOX_ASSERT(insn->next == NULL);
+      ASSERT_TRUE(insn->next == NULL);
       if (stack.empty()) {
         break;
       }
       insn = stack.back();
       stack.pop_back();
     } else {
-      SANDBOX_ASSERT(insn->next != NULL);
+      ASSERT_TRUE(insn->next != NULL);
       insn = insn->next;
     }
   }
-  SANDBOX_ASSERT(target_instructions.size() == branch_targets.size());
+  ASSERT_TRUE(target_instructions.size() == branch_targets.size());
 
   // We can now subtract the set of the branch targets from the set of all
   // instructions. This gives us a set with the instructions that nobody
@@ -312,12 +291,10 @@ void FindBranchTargets(CodeGenUnittestHelper* codegen, Instruction* prg, int) {
   for (Instructions::const_iterator iter = non_target_instructions.begin();
        iter != non_target_instructions.end();
        ++iter) {
-    SANDBOX_ASSERT(branch_targets.find(*iter) == end);
+    ASSERT_TRUE(branch_targets.find(*iter) == end);
   }
 }
 
-SANDBOX_TEST(CodeGen, FindBranchTargets) { ForAllPrograms(FindBranchTargets); }
-
 void CutGraphIntoBasicBlocks(CodeGenUnittestHelper* codegen,
                              Instruction* prg,
                              int) {
@@ -326,8 +303,8 @@ void CutGraphIntoBasicBlocks(CodeGenUnittestHelper* codegen,
   TargetsToBlocks all_blocks;
   BasicBlock* first_block =
       codegen->CutGraphIntoBasicBlocks(prg, branch_targets, &all_blocks);
-  SANDBOX_ASSERT(first_block != NULL);
-  SANDBOX_ASSERT(first_block->instructions.size() > 0);
+  ASSERT_TRUE(first_block != NULL);
+  ASSERT_TRUE(first_block->instructions.size() > 0);
   Instruction* first_insn = first_block->instructions[0];
 
   // Basic blocks are supposed to start with a branch target and end with
@@ -338,31 +315,27 @@ void CutGraphIntoBasicBlocks(CodeGenUnittestHelper* codegen,
        bb_iter != all_blocks.end();
        ++bb_iter) {
     BasicBlock* bb = bb_iter->second;
-    SANDBOX_ASSERT(bb != NULL);
-    SANDBOX_ASSERT(bb->instructions.size() > 0);
+    ASSERT_TRUE(bb != NULL);
+    ASSERT_TRUE(bb->instructions.size() > 0);
     Instruction* insn = bb->instructions[0];
-    SANDBOX_ASSERT(insn == first_insn ||
-                   branch_targets.find(insn) != branch_targets.end());
+    ASSERT_TRUE(insn == first_insn ||
+                branch_targets.find(insn) != branch_targets.end());
     for (Instructions::const_iterator insn_iter = bb->instructions.begin();;) {
       insn = *insn_iter;
       if (++insn_iter != bb->instructions.end()) {
-        SANDBOX_ASSERT(BPF_CLASS(insn->code) != BPF_JMP);
-        SANDBOX_ASSERT(BPF_CLASS(insn->code) != BPF_RET);
+        ASSERT_TRUE(BPF_CLASS(insn->code) != BPF_JMP);
+        ASSERT_TRUE(BPF_CLASS(insn->code) != BPF_RET);
       } else {
-        SANDBOX_ASSERT(BPF_CLASS(insn->code) == BPF_JMP ||
-                       BPF_CLASS(insn->code) == BPF_RET ||
-                       branch_targets.find(insn->next) != branch_targets.end());
+        ASSERT_TRUE(BPF_CLASS(insn->code) == BPF_JMP ||
+                    BPF_CLASS(insn->code) == BPF_RET ||
+                    branch_targets.find(insn->next) != branch_targets.end());
         break;
       }
-      SANDBOX_ASSERT(branch_targets.find(*insn_iter) == branch_targets.end());
+      ASSERT_TRUE(branch_targets.find(*insn_iter) == branch_targets.end());
     }
   }
 }
 
-SANDBOX_TEST(CodeGen, CutGraphIntoBasicBlocks) {
-  ForAllPrograms(CutGraphIntoBasicBlocks);
-}
-
 void MergeTails(CodeGenUnittestHelper* codegen, Instruction* prg, int flags) {
   BranchTargets branch_targets;
   codegen->FindBranchTargets(*prg, &branch_targets);
@@ -435,18 +408,14 @@ void MergeTails(CodeGenUnittestHelper* codegen, Instruction* prg, int flags) {
     }
     codegen->MergeTails(&all_blocks);
   }
-  SANDBOX_ASSERT(graph[0] == graph[1]);
+  ASSERT_TRUE(graph[0] == graph[1]);
   if (flags & HAS_MERGEABLE_TAILS) {
-    SANDBOX_ASSERT(edges[0] != edges[1]);
+    ASSERT_TRUE(edges[0] != edges[1]);
   } else {
-    SANDBOX_ASSERT(edges[0] == edges[1]);
+    ASSERT_TRUE(edges[0] == edges[1]);
   }
 }
 
-SANDBOX_TEST(CodeGen, MergeTails) {
-  ForAllPrograms(MergeTails);
-}
-
 void CompileAndCompare(CodeGenUnittestHelper* codegen, Instruction* prg, int) {
   // TopoSortBasicBlocks() has internal checks that cause it to fail, if it
   // detects a problem. Typically, if anything goes wrong, this looks to the
@@ -499,7 +468,7 @@ void CompileAndCompare(CodeGenUnittestHelper* codegen, Instruction* prg, int) {
   std::string assembly;
   std::vector<int> assembly_stack;
   for (int idx = 0; idx >= 0;) {
-    SANDBOX_ASSERT(idx < (int)bpf.size());
+    ASSERT_TRUE(idx < (int)bpf.size());
     struct sock_filter& insn = bpf[idx];
     if (BPF_CLASS(insn.code) == BPF_JMP) {
       if (BPF_OP(insn.code) == BPF_JA) {
@@ -524,11 +493,21 @@ void CompileAndCompare(CodeGenUnittestHelper* codegen, Instruction* prg, int) {
     assembly.append(reinterpret_cast<char*>(&insn.code), sizeof(insn.code));
     assembly.append(reinterpret_cast<char*>(&insn.k), sizeof(insn.k));
   }
-  SANDBOX_ASSERT(source == assembly);
+  ASSERT_TRUE(source == assembly);
 }
 
-SANDBOX_TEST(CodeGen, All) {
-  ForAllPrograms(CompileAndCompare);
-}
+const ProgramTestFunc kProgramTestFuncs[] = {
+    MakeInstruction,
+    FindBranchTargets,
+    CutGraphIntoBasicBlocks,
+    MergeTails,
+    CompileAndCompare,
+};
+
+INSTANTIATE_TEST_CASE_P(CodeGen,
+                        ProgramTest,
+                        ::testing::ValuesIn(kProgramTestFuncs));
+
+}  // namespace
 
 }  // namespace sandbox