@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();
}