Add return type tests

miland-db · Mar 20, 2024 · 3047a3c · 3047a3c
1 parent 79abfc1
commit 3047a3c
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala
@@ -74,16 +74,28 @@ class CollationStringExpressionsSuite extends QueryTest with SharedSparkSession
 
   test("Support Repeat string expression with Collation") {
     // Supported collations
-    val checks = Seq(
+    val checksRepeat = Seq(
       RepeatCollationTestCase("Spark", 2, "UTF8_BINARY", "SparkSpark"),
       RepeatCollationTestCase("Spark", 2, "UTF8_BINARY_LCASE", "SparkSpark"),
       RepeatCollationTestCase("Spark", 2, "UNICODE", "SparkSpark"),
       RepeatCollationTestCase("Spark", 2, "UNICODE_CI", "SparkSpark"),
     )
-    checks.foreach(ct => {
+    checksRepeat.foreach(ct => {
       checkAnswer(sql(s"SELECT repeat(collate('${ct.s}', '${ct.collation}'), ${ct.n})"),
         Row(ct.expectedResult))
     })
+
+    // Check return type
+    val checksReturnType = Seq(
+      RepeatCollationTestCase("Spark", 2, "UTF8_BINARY", "UTF8_BINARY"),
+      RepeatCollationTestCase("Spark", 2, "UTF8_BINARY_LCASE", "UTF8_BINARY_LCASE"),
+      RepeatCollationTestCase("Spark", 2, "UNICODE", "UNICODE"),
+      RepeatCollationTestCase("Spark", 2, "UNICODE_CI", "UNICODE_CI"),
+    )
+    checksReturnType.foreach(ct => {
+      checkAnswer(sql(s"SELECT collation(repeat(collate('${ct.s}', '${ct.collation}'), ${ct.n}))"),
+        Row(ct.expectedResult))
+    })
   }
 
   // TODO: Add more tests for other string expressions