fix test and improve performance

broadinstitute · Aug 14, 2015 · bc639f3 · bc639f3
1 parent 6c34b2a
commit bc639f3
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 23 deletions.
diff --git a/src/main/java/org/broadinstitute/hellbender/engine/dataflow/datasources/RefAPISource.java b/src/main/java/org/broadinstitute/hellbender/engine/dataflow/datasources/RefAPISource.java
@@ -4,6 +4,8 @@
 import com.google.api.services.genomics.model.ListBasesResponse;
 import com.google.api.services.genomics.model.Reference;
 import com.google.api.services.genomics.model.ReferenceSet;
+import com.google.api.services.genomics.model.SearchReferencesRequest;
+import com.google.api.services.genomics.model.SearchReferencesResponse;
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
 import com.google.cloud.genomics.dataflow.utils.GCSOptions;
 import com.google.cloud.genomics.dataflow.utils.GenomicsOptions;
@@ -67,29 +69,22 @@ private RefAPISource() { /* To prevent instantiation */ }
      * currently takes ~10 seconds, so this table should be cached and no produced per query to get the reference bases.
      * This is clumsy and should be refactored with better Genomics API use (issue 643).
      * @param pipelineOptions options needed for credentials to produce a Genomics class for the API call.
-     * @param referenceName the "name of the reference, e.g., EOSt9JOVhp3jkwE.
+     * @param referenceSetID the ID of the reference set, e.g., EOSt9JOVhp3jkwE.
      * @return returns a mapping from reference name to String id.
      */
-    public static Map<String, String> buildReferenceNameToIdTable(final PipelineOptions pipelineOptions, final String referenceName) {
+    public static Map<String, String> buildReferenceNameToIdTable(final PipelineOptions pipelineOptions, final String referenceSetID) {
         Genomics genomicsService = createGenomicsService(pipelineOptions);
-        final ReferenceSet referenceSet;
-        try {
-            referenceSet = genomicsService.referencesets().get(referenceName).execute();
-        }
-        catch ( IOException e ) {
-            throw new UserException("Could not load reference set for reference name " + referenceName, e);
-        }
-        final Map<String, String> referenceNameToIdTable = new HashMap<>();
 
+        final Map<String, String> referenceNameToIdTable = new HashMap<>();
         try {
-            for ( final String referenceID : referenceSet.getReferenceIds() ) {
-                final Reference reference = genomicsService.references().get(referenceID).execute();
-                referenceNameToIdTable.put(reference.getName(), reference.getId());
+            final SearchReferencesRequest content = new SearchReferencesRequest();
+            content.setReferenceSetId(referenceSetID);
+            final SearchReferencesResponse found = genomicsService.references().search(content).execute();
+            for (Reference r : found.getReferences()) {
+                referenceNameToIdTable.put(r.getName(), r.getId());
             }
-        }
-
-        catch ( IOException e ) {
-            throw new UserException("Error while looking up references for reference set " + referenceSet.getId(), e);
+        } catch ( IOException e ) {
+            throw new UserException("Error while looking up reference set " + referenceSetID, e);
         }
 
         return referenceNameToIdTable;

diff --git a/src/test/java/org/broadinstitute/hellbender/engine/dataflow/RefAPISourceUnitTest.java b/src/test/java/org/broadinstitute/hellbender/engine/dataflow/RefAPISourceUnitTest.java
@@ -33,10 +33,12 @@ private ReferenceBases queryReferenceAPI( final String referenceName, final Simp
         return refAPISource.getReferenceBases(p.getOptions(), refAPIMetadata, interval);
     }
 
-    @Test(groups = "cloud_todo")
+    @Test(groups = "cloud")
     public void testDummy() {
         String referenceName = "EOSt9JOVhp3jkwE";
-        SimpleInterval interval = new SimpleInterval("1", 50001, 10050000);
+        final String expected = "AAACAGGTTA";
+        // -1 because we're using closed intervals
+        SimpleInterval interval = new SimpleInterval("1", 50001, 50001 + expected.length() - 1);
         Logger logger = LogManager.getLogger(RefAPISourceUnitTest.class);
 
         GenomicsOptions options = PipelineOptionsFactory.create().as(GenomicsOptions.class);
@@ -54,12 +56,13 @@ public void testDummy() {
         ReferenceBases bases = refAPISource.getReferenceBases(p.getOptions(), refAPIMetadata, interval);
         logger.info("e");
 
-        Assert.assertEquals(new String(bases.getBases()), "AAACAGGTTA", "Wrong bases returned");
+        final String actual = new String(bases.getBases());
+        Assert.assertEquals(actual, expected, "Wrong bases returned");
         p.run();
 
     }
 
-    @Test(groups = "cloud_todo")
+    @Test(groups = "cloud")
     public void testReferenceSourceQuery() {
         final ReferenceBases bases = queryReferenceAPI("EOSt9JOVhp3jkwE", new SimpleInterval("1", 50000, 50009));
 
@@ -69,12 +72,12 @@ public void testReferenceSourceQuery() {
         Assert.assertEquals(new String(bases.getBases()), "TAAACAGGTT", "Wrong bases returned");
     }
 
-    @Test(groups = "cloud_todo", expectedExceptions = UserException.class)
+    @Test(groups = "cloud", expectedExceptions = UserException.class)
     public void testReferenceSourceQueryWithInvalidContig() {
         final ReferenceBases bases = queryReferenceAPI("EOSt9JOVhp3jkwE", new SimpleInterval("FOOCONTIG", 1, 2));
     }
 
-    @Test(groups = "cloud_todo", expectedExceptions = UserException.class)
+    @Test(groups = "cloud", expectedExceptions = UserException.class)
     public void testReferenceSourceQueryWithInvalidPosition() {
         final ReferenceBases bases = queryReferenceAPI("EOSt9JOVhp3jkwE", new SimpleInterval("1", 1000000000, 2000000000));
     }