diff --git a/test/run/test_stages.py b/test/run/test_stages.py
index 17c2db5b8..dd6617f8c 100644
--- a/test/run/test_stages.py
+++ b/test/run/test_stages.py
@@ -223,12 +223,12 @@ def run_stage_diff_test(self, test_dir: str):
 
             out_a = stack.enter_context(tempfile.TemporaryDirectory(dir="/var/tmp"))
             _ = osb.compile_file(os.path.join(test_dir, "a.json"),
-                                 checkpoints=["tree"],
+                                 checkpoints=["build", "tree"],
                                  exports=["tree"], output_dir=out_a)
 
             out_b = stack.enter_context(tempfile.TemporaryDirectory(dir="/var/tmp"))
             res = osb.compile_file(os.path.join(test_dir, "b.json"),
-                                   checkpoints=["tree"],
+                                   checkpoints=["build", "tree"],
                                    exports=["tree"], output_dir=out_b)
 
             tree1 = os.path.join(out_a, "tree")