From fb191eb8744175b6098f0775febb09aba480b262 Mon Sep 17 00:00:00 2001 From: Chen Dai Date: Wed, 20 Sep 2023 08:44:57 -0700 Subject: [PATCH] Add assertion on DF schema in IT Signed-off-by: Chen Dai --- .../spark/function/TumbleFunctionSuite.scala | 2 +- .../FlintSparkWindowingFunctionITSuite.scala | 16 +++++++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/function/TumbleFunctionSuite.scala b/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/function/TumbleFunctionSuite.scala index c6a19d92f..d79dd8f4c 100644 --- a/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/function/TumbleFunctionSuite.scala +++ b/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/function/TumbleFunctionSuite.scala @@ -23,7 +23,7 @@ class TumbleFunctionSuite extends FlintSuite with Matchers { } } - test("should fail if only argument type is wrong") { + test("should fail if argument type is wrong") { assertThrows[AnalysisException] { TumbleFunction.functionBuilder(Seq(col("timestamp").expr, lit(10).expr)) } diff --git a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkWindowingFunctionITSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkWindowingFunctionITSuite.scala index 8a32b8112..9f38958c3 100644 --- a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkWindowingFunctionITSuite.scala +++ b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkWindowingFunctionITSuite.scala @@ -7,8 +7,11 @@ package org.opensearch.flint.spark import java.sql.Timestamp +import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper + import org.apache.spark.FlintSuite import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.types.{StructField, StructType, TimestampType} class FlintSparkWindowingFunctionITSuite extends QueryTest with FlintSuite { @@ -22,12 +25,15 @@ class FlintSparkWindowingFunctionITSuite extends QueryTest with FlintSuite { .toDF("id", "timestamp") val resultDF = inputDF.selectExpr("TUMBLE(timestamp, '10 minutes') AS window") - val expectedData = Seq( - Row(Row(timestamp("2023-01-01 00:00:00"), timestamp("2023-01-01 00:10:00"))), - Row(Row(timestamp("2023-01-01 00:00:00"), timestamp("2023-01-01 00:10:00"))), - Row(Row(timestamp("2023-01-01 00:10:00"), timestamp("2023-01-01 00:20:00")))) - checkAnswer(resultDF, expectedData) + resultDF.schema shouldBe StructType.fromDDL( + "window struct NOT NULL") + checkAnswer( + resultDF, + Seq( + Row(Row(timestamp("2023-01-01 00:00:00"), timestamp("2023-01-01 00:10:00"))), + Row(Row(timestamp("2023-01-01 00:00:00"), timestamp("2023-01-01 00:10:00"))), + Row(Row(timestamp("2023-01-01 00:10:00"), timestamp("2023-01-01 00:20:00"))))) } private def timestamp(ts: String): Timestamp = Timestamp.valueOf(ts)