update RoutineTest.java

stephaniewang526 · stephaniewang526 · commit 6aac2ed57839 · 2021-05-04T11:04:25.000-04:00
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java
@@ -38,9 +38,11 @@
 public class RoutineTest {
 
   private static final RoutineId ROUTINE_ID = RoutineId.of("dataset", "routine");
+  private static final RoutineId ROUTINE_ID_TVF = RoutineId.of("dataset", "tvf_routine");
   private static final String DETERMINISM_LEVEL = "DETERMINISTIC";
   private static final String ETAG = "etag";
   private static final String ROUTINE_TYPE = "SCALAR_FUNCTION";
+  private static final String ROUTINE_TYPE_TVF = "TABLE_VALUED_FUNCTION";
   private static final Long CREATION_TIME = 10L;
   private static final Long LAST_MODIFIED_TIME = 20L;
   private static final String LANGUAGE = "SQL";
@@ -56,6 +58,18 @@ public class RoutineTest {
   private static final StandardSQLDataType RETURN_TYPE =
       StandardSQLDataType.newBuilder("FLOAT64").build();
 
+  private static final StandardSQLField COLUMN_1 =
+      StandardSQLField.newBuilder("COLUMN_1", StandardSQLDataType.newBuilder("STRING").build())
+          .build();
+  private static final StandardSQLField COLUMN_2 =
+      StandardSQLField.newBuilder("COLUMN_2", StandardSQLDataType.newBuilder("FLOAT64").build())
+          .build();
+
+  private static final List<StandardSQLField> COLUMN_LIST = ImmutableList.of(COLUMN_1, COLUMN_2);
+
+  private static final StandardSQLTableType RETURN_TABLE_TYPE =
+      StandardSQLTableType.newBuilder(COLUMN_LIST).build();
+
   private static final List<String> IMPORTED_LIBRARIES =
       ImmutableList.of("gs://foo", "gs://bar", "gs://baz");
 
@@ -75,11 +89,19 @@ public class RoutineTest {
           .setBody(BODY)
           .build();
 
+  private static final RoutineInfo ROUTINE_INFO_TVF =
+      RoutineInfo.newBuilder(ROUTINE_ID_TVF)
+          .setBody(BODY)
+          .setRoutineType(ROUTINE_TYPE_TVF)
+          .setReturnTableType(RETURN_TABLE_TYPE)
+          .build();
+
   @Rule public MockitoRule rule;
 
   private BigQuery bigquery;
   private BigQueryOptions mockOptions;
   private Routine expectedRoutine;
+  private Routine expectedRoutineTvf;
   private Routine routine;
 
   @Before
@@ -88,6 +110,7 @@ public void setUp() {
     mockOptions = mock(BigQueryOptions.class);
     when(bigquery.getOptions()).thenReturn(mockOptions);
     expectedRoutine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
+    expectedRoutineTvf = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO_TVF));
     routine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
   }
 
@@ -114,6 +137,7 @@ public void testBuilder() {
   @Test
   public void testToBuilder() {
     compareRoutineInfo(expectedRoutine, expectedRoutine.toBuilder().build());
+    compareRoutineInfo(expectedRoutineTvf, expectedRoutineTvf.toBuilder().build());
   }
 
   @Test
@@ -200,6 +224,7 @@ public void compareRoutineInfo(RoutineInfo expected, RoutineInfo value) {
     assertEquals(expected.getLanguage(), value.getLanguage());
     assertEquals(expected.getArguments(), value.getArguments());
     assertEquals(expected.getReturnType(), value.getReturnType());
+    assertEquals(expected.getReturnTableType(), value.getReturnTableType());
     assertEquals(expected.getImportedLibraries(), value.getImportedLibraries());
     assertEquals(expected.getBody(), value.getBody());
     assertEquals(expected.hashCode(), value.hashCode());
