Add a test case for AggregatedDialect.isCascadingTruncateTable.

viirya · viirya · commit d4fadbbab30c · 2017-09-20T06:47:51.000Z
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala
@@ -46,7 +46,7 @@ private class AggregatedDialect(dialects: List[JdbcDialect]) extends JdbcDialect
     // If any dialect claims cascading truncate, this dialect is also cascading truncate.
     // Otherwise, if any dialect has unknown cascading truncate, this dialect is also unknown.
     val cascading = dialects.flatMap(_.isCascadingTruncateTable()).reduceOption(_ || _)
-    if (cascading.get) {
+    if (cascading.getOrElse(false)) {
       cascading
     } else {
       if (dialects.exists(_.isCascadingTruncateTable().isEmpty)) {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala
@@ -747,19 +747,34 @@ class JDBCSuite extends SparkFunSuite
     assert(agg.getCatalystType(0, "", 1, null) === Some(LongType))
     assert(agg.getCatalystType(1, "", 1, null) === Some(StringType))
     assert(agg.isCascadingTruncateTable() === Some(true))
+  }
 
-    val agg2 = new AggregatedDialect(List(new JdbcDialect {
-      override def canHandle(url: String) : Boolean = url.startsWith("jdbc:h2:")
+  test("Aggregated dialects: isCascadingTruncateTable") {
+    def genDialect(cascadingTruncateTable: Option[Boolean]): JdbcDialect = new JdbcDialect {
+      override def canHandle(url: String): Boolean = true
       override def getCatalystType(
-          sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] =
-        if (sqlType % 2 == 0) {
-          Some(LongType)
-        } else {
-          None
-        }
-      override def isCascadingTruncateTable(): Option[Boolean] = Some(false)
-    }, testH2Dialect))
-    assert(agg2.isCascadingTruncateTable() === None)
+        sqlType: Int,
+        typeName: String,
+        size: Int,
+        md: MetadataBuilder): Option[DataType] = None
+      override def isCascadingTruncateTable(): Option[Boolean] = cascadingTruncateTable
+    }
+
+    val dialectCombination = Seq(
+      List(genDialect(Some(true)), genDialect(Some(false)), genDialect(None)),
+      List(genDialect(Some(true)), genDialect(Some(true)), genDialect(None)),
+      List(genDialect(Some(false)), genDialect(Some(false)), genDialect(None)),
+      List(genDialect(Some(true)), genDialect(Some(true))),
+      List(genDialect(Some(false)), genDialect(Some(false))),
+      List(genDialect(None), genDialect(None))
+    )
+
+    val expectedCascading = Seq(Some(true), Some(true), None, Some(true), Some(false), None)
+
+    dialectCombination.zip(expectedCascading).foreach { case (dialects, cascading) =>
+      val agg = new AggregatedDialect(dialects)
+      assert(agg.isCascadingTruncateTable() === cascading)
+    }
   }
 
   test("DB2Dialect type mapping") {
