chore: Improve test assertions in plan stability suite (apache#2505)

Author: andygrove

spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala

@@ -89,79 +89,64 @@ trait CometPlanStabilitySuite extends DisableAdaptiveExecutionSuite with TPCDSBa
     new File(goldenFilePath, goldenFileName)
   }
 
-  private def isApproved(
-      dir: File,
-      actualSimplifiedPlan: String,
-      actualExplain: String): Boolean = {
-    val simplifiedFile = new File(dir, "simplified.txt")
-    val expectedSimplified = FileUtils.readFileToString(simplifiedFile, StandardCharsets.UTF_8)
-    val explainFile = new File(dir, "explain.txt")
-    val expectedExplain = FileUtils.readFileToString(explainFile, StandardCharsets.UTF_8)
-    expectedSimplified == actualSimplifiedPlan && expectedExplain == actualExplain
-  }
-
   /**
    * Serialize and save this SparkPlan. The resulting file is used by [[checkWithApproved]] to
    * check stability.
    *
-   * @param plan
-   *   the SparkPlan
-   * @param name
-   *   the name of the query
+   * @param dir
+   *   the directory to write to
+   * @param simplified
+   *   the simplified plan
    * @param explain
    *   the full explain output; this is saved to help debug later as the simplified plan is not
    *   too useful for debugging
    */
-  private def generateGoldenFile(plan: SparkPlan, name: String, explain: String): Unit = {
-    val dir = getDirForTest(name)
-    val simplified = getSimplifiedPlan(plan)
-    val foundMatch = dir.exists() && isApproved(dir, simplified, explain)
-
-    if (!foundMatch) {
-      FileUtils.deleteDirectory(dir)
-      if (!dir.mkdirs()) {
-        fail(s"Could not create dir: $dir")
-      }
-
-      val file = new File(dir, "simplified.txt")
-      FileUtils.writeStringToFile(file, simplified, StandardCharsets.UTF_8)
-      val fileOriginalPlan = new File(dir, "explain.txt")
-      FileUtils.writeStringToFile(fileOriginalPlan, explain, StandardCharsets.UTF_8)
-      logDebug(s"APPROVED: $file $fileOriginalPlan")
+  private def generateGoldenFile(dir: File, simplified: String, explain: String): Unit = {
+    FileUtils.deleteDirectory(dir)
+    if (!dir.mkdirs()) {
+      fail(s"Could not create dir: $dir")
     }
+
+    val file = new File(dir, "simplified.txt")
+    FileUtils.writeStringToFile(file, simplified, StandardCharsets.UTF_8)
+    val fileOriginalPlan = new File(dir, "explain.txt")
+    FileUtils.writeStringToFile(fileOriginalPlan, explain, StandardCharsets.UTF_8)
+    logDebug(s"APPROVED: $file $fileOriginalPlan")
   }
 
-  private def checkWithApproved(plan: SparkPlan, name: String, explain: String): Unit = {
-    val dir = getDirForTest(name)
-    val tempDir = FileUtils.getTempDirectory
-    val actualSimplified = getSimplifiedPlan(plan)
-    val foundMatch = isApproved(dir, actualSimplified, explain)
+  private def checkWithApproved(
+      dir: File,
+      name: String,
+      simplified: String,
+      explain: String): Unit = {
 
-    if (!foundMatch) {
-      // show diff with last approved
-      val approvedSimplifiedFile = new File(dir, "simplified.txt")
-      val approvedExplainFile = new File(dir, "explain.txt")
+    val approvedSimplifiedFile = new File(dir, "simplified.txt")
+    val approvedExplainFile = new File(dir, "explain.txt")
 
-      val actualSimplifiedFile = new File(tempDir, s"$name.actual.simplified.txt")
-      val actualExplainFile = new File(tempDir, s"$name.actual.explain.txt")
+    // write actual files out for debugging
+    val tempDir = FileUtils.getTempDirectory
+    val actualSimplifiedFile = new File(tempDir, s"$name.actual.simplified.txt")
+    val actualExplainFile = new File(tempDir, s"$name.actual.explain.txt")
+    FileUtils.writeStringToFile(actualSimplifiedFile, simplified, StandardCharsets.UTF_8)
+    FileUtils.writeStringToFile(actualExplainFile, explain, StandardCharsets.UTF_8)
 
-      val approvedSimplified =
-        FileUtils.readFileToString(approvedSimplifiedFile, StandardCharsets.UTF_8)
-      // write out for debugging
-      FileUtils.writeStringToFile(actualSimplifiedFile, actualSimplified, StandardCharsets.UTF_8)
-      FileUtils.writeStringToFile(actualExplainFile, explain, StandardCharsets.UTF_8)
+    comparePlans("simplified", approvedSimplifiedFile, actualSimplifiedFile)
+    comparePlans("explain", approvedExplainFile, actualExplainFile)
+  }
 
+  private def comparePlans(planType: String, expectedFile: File, actualFile: File): Unit = {
+    val expected = FileUtils.readFileToString(expectedFile, StandardCharsets.UTF_8)
+    val actual = FileUtils.readFileToString(actualFile, StandardCharsets.UTF_8)
+    if (expected != actual) {
       fail(s"""
-           |Plans did not match:
-           |last approved simplified plan: ${approvedSimplifiedFile.getAbsolutePath}
-           |last approved explain plan: ${approvedExplainFile.getAbsolutePath}
-           |
-           |$approvedSimplified
-           |
-           |actual simplified plan: ${actualSimplifiedFile.getAbsolutePath}
-           |actual explain plan: ${actualExplainFile.getAbsolutePath}
-           |
-           |$actualSimplified
+              |Plans did not match:
+              |last approved $planType plan: ${expectedFile.getAbsolutePath}
+              |
+              |$expected
+              |
+              |actual $planType plan: ${actualFile.getAbsolutePath}
+              |
+              |$actual
         """.stripMargin)
     }
   }
@@ -274,13 +259,16 @@ trait CometPlanStabilitySuite extends DisableAdaptiveExecutionSuite with TPCDSBa
       val qe = sql(queryString).queryExecution
       val plan = qe.executedPlan
       val explain = normalizeLocation(normalizeIds(qe.explainString(FormattedMode)))
-
+      val simplified = getSimplifiedPlan(plan)
       assert(ValidateRequirements.validate(plan))
 
+      val name = query + suffix
+      val dir = getDirForTest(name)
+
       if (regenerateGoldenFiles) {
-        generateGoldenFile(plan, query + suffix, explain)
+        generateGoldenFile(dir, simplified, explain)
       } else {
-        checkWithApproved(plan, query + suffix, explain)
+        checkWithApproved(dir, name, simplified, explain)
       }
     }
   }