Fix test, don't use RowEncoder because it doesn't transfer metadata

AbsaOSS · Feb 8, 2022 · c9d1b22 · c9d1b22
1 parent e65ac09
commit c9d1b22
Showing 1 changed file with 22 additions and 3 deletions.
diff --git a/...stor/implementation/transformer/avro/confluent/TestConfluentAvroEncodingTransformer.scala b/...stor/implementation/transformer/avro/confluent/TestConfluentAvroEncodingTransformer.scala
@@ -19,7 +19,8 @@ import io.confluent.kafka.schemaregistry.client.MockSchemaRegistryClient
 import org.apache.commons.configuration2.BaseConfiguration
 import org.apache.commons.configuration2.convert.DefaultListDelimiterHandler
 import org.apache.spark.sql.Row
-import org.apache.spark.sql.catalyst.encoders.RowEncoder
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
+import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.execution.streaming.MemoryStream
 import org.apache.spark.sql.streaming.Trigger
 import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType}
@@ -36,6 +37,8 @@ import za.co.absa.hyperdrive.ingestor.implementation.transformer.avro.confluent.
 import za.co.absa.hyperdrive.ingestor.implementation.utils.AbrisConfigUtil
 import za.co.absa.hyperdrive.ingestor.implementation.writer.kafka.KafkaStreamWriter
 
+import scala.reflect.ClassTag
+
 class TestConfluentAvroEncodingTransformer extends FlatSpec with Matchers with BeforeAndAfter with SparkTestBase {
 
   private val topic = "topic"
@@ -173,7 +176,9 @@ class TestConfluentAvroEncodingTransformer extends FlatSpec with Matchers with B
       StructField("col1", IntegerType, nullable = true,
         new MetadataBuilder().putString(SparkMetadataKeys.DefaultValueKey, "42").build()))
     )
-    val memoryStream = new MemoryStream[Row](1, spark.sqlContext)(RowEncoder(schema))
+
+    import scala.collection.JavaConverters._
+    val df = spark.createDataFrame(Seq[Row]().asJava, schema)
 
     val config = new BaseConfiguration()
     config.setListDelimiterHandler(new DefaultListDelimiterHandler(','))
@@ -197,10 +202,24 @@ class TestConfluentAvroEncodingTransformer extends FlatSpec with Matchers with B
     }
 
     // when
-    encoder.transform(memoryStream.toDF())
+    encoder.transform(df)
 
     // then
     val schemaMetadata = mockSchemaRegistryClient.getLatestSchemaMetadata(s"$topic-value")
     schemaMetadata.getSchema shouldBe expectedSchemaString
   }
 }
+
+class MetadataSafeExpressionEncoder[T](objSerializer2: Expression, objDeserializer2: Expression, clsTag2: ClassTag[T])
+  extends ExpressionEncoder[T](objSerializer2, objDeserializer2, clsTag2) {
+
+  override val schema: StructType = StructType(serializer.map { s =>
+    StructField(s.name, s.dataType, s.nullable, s.metadata)
+  })
+
+  override protected val schemaString: String =
+    schema
+      .zip(attrs)
+      .map { case(f, a) => s"${f.name}$a: ${f.dataType.simpleString}"}.mkString(", ")
+
+}