From bd92170a3152eaa34c51d555bdf527b9ec9dcf41 Mon Sep 17 00:00:00 2001
From: Sabir <sabir.akhadov@databricks.com>
Date: Tue, 21 May 2024 23:27:29 +0200
Subject: [PATCH] .

---
 .../columnmapping/RemoveColumnMappingSuiteUtils.scala      | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/spark/src/test/scala/org/apache/spark/sql/delta/columnmapping/RemoveColumnMappingSuiteUtils.scala b/spark/src/test/scala/org/apache/spark/sql/delta/columnmapping/RemoveColumnMappingSuiteUtils.scala
index 420435855cf..8ec42ff2d47 100644
--- a/spark/src/test/scala/org/apache/spark/sql/delta/columnmapping/RemoveColumnMappingSuiteUtils.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/delta/columnmapping/RemoveColumnMappingSuiteUtils.scala
@@ -102,8 +102,11 @@ trait RemoveColumnMappingSuiteUtils extends QueryTest with DeltaColumnMappingSui
       spark.table(tableName = testTableName).select(logicalColumnName),
       originalData)
     val newSnapshot = deltaLog.update()
-    // Drop feature adds 2 empty commits.
-    val versionsAddedByRewrite = if (droppedFeature) 3 else 1
+    val versionsAddedByRewrite = if (droppedFeature) {
+      2
+    } else {
+      1
+    }
     assert(newSnapshot.version - startingVersion == versionsAddedByRewrite,
       s"Should rewrite the table in $versionsAddedByRewrite commits.")