Use a handmade TestReferenceMultiSource in tests instead of a mock.

src/main/java/org/broadinstitute/hellbender/engine/datasources/ReferenceMultiSource.java

@@ -1,6 +1,7 @@
 package org.broadinstitute.hellbender.engine.datasources;
 
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
+import com.google.common.annotations.VisibleForTesting;
 import org.broadinstitute.hellbender.utils.SerializableFunction;
 import htsjdk.samtools.SAMSequenceDictionary;
 import htsjdk.samtools.reference.ReferenceSequenceFileFactory;
@@ -20,14 +21,17 @@
 /**
  * Wrapper to load a reference sequence from the Google Genomics API, or a file stored on HDFS or locally.
  *
- * This class needs to be mocked, so it cannot be declared final.
+ * This class needs to subclassed by test code, so it cannot be declared final.
  */
 public class ReferenceMultiSource implements ReferenceSource, Serializable {
     private static final long serialVersionUID = 1L;
 
     private ReferenceSource referenceSource;
     private SerializableFunction<GATKRead, SimpleInterval> referenceWindowFunction;
 
+    @VisibleForTesting
+    protected ReferenceMultiSource() {};
+
     /**
      * @param pipelineOptions the pipeline options; must be GCSOptions if using the Google Genomics API
      * @param referenceURL the name of the reference (if using the Google Genomics API), or a path to the reference file

src/test/java/org/broadinstitute/hellbender/engine/spark/AddContextDataToReadSparkUnitTest.java

@@ -1,6 +1,7 @@
 package org.broadinstitute.hellbender.engine.spark;
 
 import com.google.api.services.genomics.model.Read;
+import com.google.cloud.dataflow.sdk.options.PipelineOptions;
 import com.google.cloud.dataflow.sdk.values.KV;
 import com.google.common.collect.Lists;
 import htsjdk.samtools.SAMRecord;
@@ -13,14 +14,13 @@
 import org.broadinstitute.hellbender.engine.ReadContextData;
 import org.broadinstitute.hellbender.engine.datasources.ReferenceWindowFunctions;
 import org.broadinstitute.hellbender.engine.datasources.ReferenceMultiSource;
+import org.broadinstitute.hellbender.utils.SerializableFunction;
 import org.broadinstitute.hellbender.utils.SimpleInterval;
 import org.broadinstitute.hellbender.utils.read.GATKRead;
 import org.broadinstitute.hellbender.utils.reference.ReferenceBases;
 import org.broadinstitute.hellbender.utils.test.BaseTest;
 import org.broadinstitute.hellbender.utils.test.FakeReferenceSource;
 import org.broadinstitute.hellbender.utils.variant.GATKVariant;
-import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
 import org.testng.Assert;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -31,8 +31,6 @@
 import java.util.List;
 import java.util.Map;
 
-import static org.mockito.Mockito.*;
-
 public class AddContextDataToReadSparkUnitTest extends BaseTest {
     @DataProvider(name = "bases")
     public Object[][] bases() {
@@ -62,14 +60,12 @@ public void addContextDataTest(List<GATKRead> reads, List<GATKVariant> variantLi
         JavaRDD<GATKRead> rddReads = ctx.parallelize(reads);
         JavaRDD<GATKVariant> rddVariants = ctx.parallelize(variantList);
 
-        ReferenceMultiSource mockSource = mock(ReferenceMultiSource.class, withSettings().serializable());
-        when(mockSource.getReferenceBases(isNull(), any(SimpleInterval.class))).then(new ReferenceBasesAnswer());
-        when(mockSource.getReferenceWindowFunction()).thenReturn(ReferenceWindowFunctions.IDENTITY_FUNCTION);
         SAMSequenceDictionary sd = new SAMSequenceDictionary(Lists.newArrayList(new SAMSequenceRecord("1", 100000), new SAMSequenceRecord("2", 100000)));
-        when(mockSource.getReferenceSequenceDictionary(null)).thenReturn(sd);
-
-        JavaPairRDD<GATKRead, ReadContextData> rddActual = AddContextDataToReadSpark.add(ctx, rddReads, mockSource, rddVariants, null, joinStrategy,
+        JavaPairRDD<GATKRead, ReadContextData> rddActual = AddContextDataToReadSpark.add(ctx, rddReads,
+                new TestMultiReferenceSource(sd),
+                rddVariants, null, joinStrategy,
                 sd, 10000, 1000);
+
         Map<GATKRead, ReadContextData> actual = rddActual.collectAsMap();
 
         Assert.assertEquals(actual.size(), expectedReadContextData.size());
@@ -84,11 +80,38 @@ public void addContextDataTest(List<GATKRead> reads, List<GATKVariant> variantLi
         }
     }
 
-    static class ReferenceBasesAnswer implements Answer<ReferenceBases>, Serializable {
+    // Provide a fake implementation of this class for testing. We can't use a real mock since this is used as a Spark
+    // broadcast variable. Mocks are mutated when they're accessed, which can result in ConcurrentModificationExceptions
+    // during serialization/broadcast.
+    static class TestMultiReferenceSource extends ReferenceMultiSource implements Serializable {
         private static final long serialVersionUID = 1L;
+
+        final SAMSequenceDictionary sequenceDictionary;
+
+        public TestMultiReferenceSource(final SAMSequenceDictionary sd) {
+            sequenceDictionary = sd;
+        }
+
+        @Override
+        public ReferenceBases getReferenceBases(final PipelineOptions pipelineOptions, final SimpleInterval interval) throws IOException {
+            return FakeReferenceSource.bases(interval);
+        }
+
         @Override
-        public ReferenceBases answer(InvocationOnMock invocation) throws Throwable {
-            return FakeReferenceSource.bases(invocation.getArgument(1));
+        public SAMSequenceDictionary getReferenceSequenceDictionary(final SAMSequenceDictionary optReadSequenceDictionaryToMatch) {
+            return sequenceDictionary;
         }
+
+        @Override
+        public boolean isCompatibleWithSparkBroadcast(){
+            return true;
+        }
+
+        @Override
+        public SerializableFunction<GATKRead, SimpleInterval> getReferenceWindowFunction() {
+            return ReferenceWindowFunctions.IDENTITY_FUNCTION;
+        }
+
     }
+
 }