ReferenceMetadataInputReader implements InputCitationReader<String, ReferenceMetadata>, Serializable { @Override public JavaPairRDD<String, ReferenceMetadata> readCitations(JavaSparkContext sparkContext, String inputCitationsPath) { JavaRDD<DocumentMetadata> fullDocuments = avroLoader.loadJavaRDD(sparkContext, inputCitationsPath, DocumentMetadata.class); JavaPairRDD<String, ReferenceMetadata> references = fullDocuments .flatMapToPair(fullDocument -> fullDocument.getReferences().stream() .map(reference -> new Tuple2<String, ReferenceMetadata>( buildCitationId(fullDocument.getId().toString(), reference), reference)) .collect(Collectors.toList()) ); return references; }  @Override JavaPairRDD<String, ReferenceMetadata> readCitations(JavaSparkContext sparkContext, String inputCitationsPath);  }
@SuppressWarnings("unchecked") @Test public void readCitations() throws Exception { doReturn(inputRecords).when(sparkContext).newAPIHadoopFile(any(), any(), any(), any(), any()); doReturn(documents).when(inputRecords).map(any()); doReturn(citations).when(documents).flatMapToPair(any()); JavaPairRDD<String, ReferenceMetadata> retCitations = referenceMetadataInputReader.readCitations(sparkContext, "/some/path"); assertTrue(retCitations == citations); verify(sparkContext).newAPIHadoopFile( eq("/some/path"), eq(AvroKeyInputFormat.class), eq(DocumentMetadata.class), eq(NullWritable.class), isA(Configuration.class)); verify(documents).flatMapToPair(documentToCitationsFunction.capture()); assertDocToCitationsFunction(documentToCitationsFunction.getValue()); }