@Test(dataProvider = "bases") public void addContextDataTest( List<GATKRead> reads, List<Variant> variantList, List<KV<GATKRead, ReferenceBases>> kvReadRefBases, List<KV<GATKRead, ReadContextData>> kvReadContextData, List<SimpleInterval> intervals, List<KV<GATKRead, Iterable<Variant>>> kvReadiVariant) { Pipeline p = GATKTestPipeline.create(); DataflowUtils.registerGATKCoders(p); PCollection<GATKRead> pReads = DataflowTestUtils.pCollectionCreateAndVerify(p, reads, new GATKReadCoder()); PCollection<KV<GATKRead, ReferenceBases>> pReadRef = DataflowTestUtils.pCollectionCreateAndVerify( p, kvReadRefBases, KvCoder.of(new GATKReadCoder(), SerializableCoder.of(ReferenceBases.class))); PCollection<KV<GATKRead, Iterable<Variant>>> pReadVariants = p.apply( Create.of(kvReadiVariant) .withCoder(KvCoder.of(new GATKReadCoder(), IterableCoder.of(new VariantCoder())))); PCollection<KV<GATKRead, ReadContextData>> joinedResults = AddContextDataToRead.join(pReads, pReadRef, pReadVariants); PCollection<KV<GATKRead, ReadContextData>> pkvReadContextData = p.apply( Create.of(kvReadContextData) .withCoder(KvCoder.of(new GATKReadCoder(), new ReadContextDataCoder()))); DataflowTestUtils.keyReadContextDataMatcher(joinedResults, pkvReadContextData); p.run(); }
@Test(dataProvider = "bases") public void fullTest( List<GATKRead> reads, List<Variant> variantList, List<KV<GATKRead, ReferenceBases>> kvReadRefBases, List<KV<GATKRead, ReadContextData>> kvReadContextData, List<SimpleInterval> intervals, List<KV<GATKRead, Iterable<Variant>>> kvReadiVariant) { Pipeline p = GATKTestPipeline.create(); DataflowUtils.registerGATKCoders(p); PCollection<GATKRead> pReads = DataflowTestUtils.pCollectionCreateAndVerify(p, reads, new GATKReadCoder()); PCollection<Variant> pVariant = p.apply(Create.of(variantList)); VariantsDataflowSource mockVariantsSource = mock(VariantsDataflowSource.class); when(mockVariantsSource.getAllVariants()).thenReturn(pVariant); RefAPISource mockSource = mock(RefAPISource.class, withSettings().serializable()); for (SimpleInterval i : intervals) { when(mockSource.getReferenceBases( any(PipelineOptions.class), any(RefAPIMetadata.class), eq(i))) .thenReturn(FakeReferenceSource.bases(i)); } String referenceName = "refName"; String refId = "0xbjfjd23f"; Map<String, String> referenceNameToIdTable = Maps.newHashMap(); referenceNameToIdTable.put(referenceName, refId); RefAPIMetadata refAPIMetadata = new RefAPIMetadata(referenceName, referenceNameToIdTable); RefAPISource.setRefAPISource(mockSource); PCollection<KV<GATKRead, ReadContextData>> result = AddContextDataToRead.add(pReads, /*mockSource,*/ refAPIMetadata, mockVariantsSource); PCollection<KV<GATKRead, ReadContextData>> pkvReadContextData = p.apply( Create.of(kvReadContextData) .withCoder(KvCoder.of(new GATKReadCoder(), new ReadContextDataCoder()))); DataflowTestUtils.keyReadContextDataMatcher(result, pkvReadContextData); p.run(); }