PLUGIN-296 Add integration test for SpannerSource getSchema from ImportQuery

Author: ardian4

integration-test-remote/src/test/java/io/cdap/cdap/app/etl/gcp/GoogleCloudSpannerTest.java

@@ -125,7 +125,11 @@ public class GoogleCloudSpannerTest extends DataprocETLTestBase {
     "schema",
     Schema.Field.of("ID", Schema.nullableOf(Schema.of(Schema.Type.LONG))),
     Schema.Field.of("StringCol", Schema.nullableOf(Schema.of(Schema.Type.STRING))),
-    Schema.Field.of("BOOL_COL", Schema.nullableOf(Schema.of(Schema.Type.BOOLEAN)))
+    Schema.Field.of("BoolCol", Schema.nullableOf(Schema.of(Schema.Type.BOOLEAN))),
+    Schema.Field.of("TimestampCol", Schema.nullableOf(Schema.of(Schema.LogicalType.TIMESTAMP_MICROS))),
+    Schema.Field.of("ArrayIntCol", Schema.arrayOf(Schema.nullableOf(Schema.of(Schema.Type.LONG)))),
+    Schema.Field.of("BytesCol", Schema.nullableOf(Schema.of(Schema.Type.BYTES))),
+    Schema.Field.of("DateCol", Schema.nullableOf(Schema.of(Schema.LogicalType.DATE)))
   );
 
   private static final ZonedDateTime NOW = ZonedDateTime.now();
@@ -168,6 +172,10 @@ public class GoogleCloudSpannerTest extends DataprocETLTestBase {
       .set("ID").to(3)
       .set("STRING_COL").to("some string")
       .set("BOOL_COL").to(false)
+      .set("TIMESTAMP_COL").to(Timestamp.ofTimeSecondsAndNanos(NOW.toEpochSecond(), NOW.getNano()))
+      .set("ARRAY_INT_COL").toInt64Array(Arrays.asList(1L, 2L, null))
+      .set("BYTES_COL").to(ByteArray.copyFrom("some value".getBytes()))
+      .set("DATE_COL").to(Date.fromYearMonthDay(NOW.getYear(), NOW.getMonthValue(), NOW.getDayOfMonth()))
     .build()
   );
 
@@ -331,7 +339,7 @@ private void testReadWithImportQuery(Engine engine) throws Exception {
       .put("instance", "${instance}")
       .put("database", "${database}")
       .put("table", "${srcTable}")
-      .put("importQuery","Select ID, STRING_COL as StringCol, BOOL_COL from " + SOURCE_TABLE_NAME)
+      .put("importQuery","${importQuery}")
       .build();
 
     Map<String, String> sinkProperties = new ImmutableMap.Builder<String, String>()
@@ -340,7 +348,6 @@ private void testReadWithImportQuery(Engine engine) throws Exception {
       .put("instance", "${instance}")
       .put("database", "${database}")
       .put("table", "${dstTable}")
-      .put("serviceAccountType", "JSON")
       .put("schema", IMPORT_SCHEMA.toString())
       .put("keys", "${keys}")
       .build();
@@ -356,6 +363,8 @@ private void testReadWithImportQuery(Engine engine) throws Exception {
     args.put("srcTable", SOURCE_TABLE_NAME);
     args.put("dstTable", nonExistentSinkTableName);
     args.put("keys", "ID");
+    args.put("importQuery","Select ID, STRING_COL as StringCol, BOOL_COL as BoolCol, TIMESTAMP_COL as TimestampCol, " +
+      "ARRAY_INT_COL as ArrayIntCol, BYTES_COL as BytesCol, DATE_COL as DateCol from " + SOURCE_TABLE_NAME);
     startWorkFlow(applicationManager, ProgramRunStatus.COMPLETED, args);
 
     ResultSet resultSet = spanner.getDatabaseClient(database.getId())
@@ -372,10 +381,15 @@ private void testReadWithImportQuery(Engine engine) throws Exception {
     Assert.assertEquals(secondRowExpected.keySet().size(), resultSet.getColumnCount());
     Assert.assertEquals(secondRowExpected.get("ID").getInt64(), resultSet.getLong("ID"));
     Assert.assertEquals(secondRowExpected.get("STRING_COL").getString(), resultSet.getString("StringCol"));
-    Assert.assertEquals(secondRowExpected.get("BOOL_COL").getBool(), resultSet.getBoolean("BOOL_COL"));
+    Assert.assertEquals(secondRowExpected.get("BOOL_COL").getBool(), resultSet.getBoolean("BoolCol"));
+    Assert.assertEquals(secondRowExpected.get("TIMESTAMP_COL").getTimestamp(), resultSet.getTimestamp("TimestampCol"));
+    Assert.assertEquals(secondRowExpected.get("ARRAY_INT_COL").getInt64Array(), resultSet.getLongList("ArrayIntCol"));
+    Assert.assertEquals(secondRowExpected.get("BYTES_COL").getBytes(), resultSet.getBytes("BytesCol"));
+    Assert.assertEquals(secondRowExpected.get("DATE_COL").getDate(), resultSet.getDate("DateCol"));
+    spanner.getDatabaseClient(database.getId()).singleUse()
+      .executeQuery(Statement.of(String.format("drop table %s;", nonExistentSinkTableName)));
   }
 
-
   //TODO:(CDAP-16040) re-enable once plugin is fixed
   //@Test
   public void testReadAndStoreInNewTableWithNoSourceSchema() throws Exception {