Added another test for BigQueryUtils (#6861)

broadinstitute · Nov 3, 2020 · 2d107f8 · 2d107f8
1 parent 39ad4ee
commit 2d107f8
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/src/test/java/org/broadinstitute/hellbender/utils/bigquery/BigQueryUtilsUnitTest.java b/src/test/java/org/broadinstitute/hellbender/utils/bigquery/BigQueryUtilsUnitTest.java
@@ -9,10 +9,7 @@
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
-import java.util.HashMap;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 /**
  * A class to test the functionality of {@link BigQueryUtils}.
@@ -76,13 +73,26 @@ public void testQueryWithStorageAPI() {
 
         final StorageAPIAvroReader result = BigQueryUtils.executeQueryWithStorageAPI(query, fieldsToRetrieve, BIGQUERY_TEST_PROJECT);
 
+        int rowCount = 0;
+        final Set<String> retrievedNames = new HashSet<>();
+
         while( result.hasNext() ) {
             GenericRecord row = result.next();
             String name = row.get("name").toString();
+            retrievedNames.add(name);
+            rowCount++;
 
             Assert.assertTrue(expectedNamesAndAges.containsKey(name), "Unexpected name " + name + " returned from query " + query);
             Assert.assertNull(row.get("age"), "Age was unexpectedly retrieved.");
         }
+
+        final Set<String> expectedNames = new HashSet<>();
+        expectedNames.add("Fred");
+        expectedNames.add("Matilda");
+        expectedNames.add("Harry");
+
+        Assert.assertEquals(retrievedNames, expectedNames, "Retrieved names did not match.");
+        Assert.assertEquals(rowCount, expectedNamesAndAges.size(), "Size of returned results does not match expected.");
     }
 
     private Map<String, String> getAllExpectedNamesAndAges() {