CoreCitationMatchingService { public JavaPairRDD<MatchableEntity, IdWithSimilarity> matchCitations(JavaPairRDD<String, MatchableEntity> citations, JavaPairRDD<String, MatchableEntity> documents) { JavaPairRDD<String, String> citIdDocIdPairs = matchCitDocByHashes(citations, documents); JavaPairRDD<String, MatchableEntity> citIdDocPairs = documentAttacher.attachDocuments(citIdDocIdPairs, documents); JavaPairRDD<MatchableEntity, MatchableEntity> citDocPairs = citationAttacher.attachCitationsAndLimitDocs(citIdDocPairs, citations); JavaPairRDD<MatchableEntity, IdWithSimilarity> matchedCitations = bestMatchedCitationPicker.pickBest(citDocPairs); return matchedCitations; }  JavaPairRDD<MatchableEntity, IdWithSimilarity> matchCitations(JavaPairRDD<String, MatchableEntity> citations, JavaPairRDD<String, MatchableEntity> documents); void setDocumentAttacher(DocumentAttacher documentAttacher); void setCitationAttacher(CitationAttacherWithMatchedLimiter citationAttacher); void setBestMatchedCitationPicker(BestMatchedCitationPicker bestMatchedCitationPicker); void setSparkContext(JavaSparkContext sparkContext); void setMatchableEntityHashers(List<Pair<MatchableEntityHasher, MatchableEntityHasher>> matchableEntityHashers); void setMaxHashBucketSize(long maxHashBucketSize); void setHeuristicHashCitationMatcherFactory(
            HeuristicHashCitationMatcherFactory heuristicHashCitationMatcherFactory);  }
@Test public void matchCitations() { HeuristicHashCitationMatcher heuristicHashCitationMatcher1 = mock(HeuristicHashCitationMatcher.class); HeuristicHashCitationMatcher heuristicHashCitationMatcher2 = mock(HeuristicHashCitationMatcher.class); when(heuristicHashCitationMatcherFactory.create(citationHasher1, documentHasher1, 100)).thenReturn(heuristicHashCitationMatcher1); when(heuristicHashCitationMatcherFactory.create(citationHasher2, documentHasher2, 100)).thenReturn(heuristicHashCitationMatcher2); when(heuristicHashCitationMatcher1.matchCitations(citations, documents, true)).thenReturn(new HeuristicHashMatchingResult(matched1, unmatched1)); when(heuristicHashCitationMatcher2.matchCitations(unmatched1, documents, false)).thenReturn(new HeuristicHashMatchingResult(matched2, null)); doReturn(emptyRDD).when(sparkContext).parallelizePairs(any()); doReturn(joinedMatched1).when(emptyRDD).union(matched1); doReturn(joinedMatched2).when(joinedMatched1).union(matched2); when(documentAttacher.attachDocuments(joinedMatched2, documents)).thenReturn(documentAttached); when(citationAttacher.attachCitationsAndLimitDocs(documentAttached, citations)).thenReturn(citationAttached); when(bestMatchedCitationPicker.pickBest(citationAttached)).thenReturn(matchedCitations); JavaPairRDD<MatchableEntity, IdWithSimilarity> retMatchedCitations = coreCitationMatchingService.matchCitations(citations, documents); assertTrue(retMatchedCitations == matchedCitations); verify(heuristicHashCitationMatcherFactory).create(citationHasher1, documentHasher1, 100); verify(heuristicHashCitationMatcherFactory).create(citationHasher2, documentHasher2, 100); verify(heuristicHashCitationMatcher1).matchCitations(citations, documents, true); verify(heuristicHashCitationMatcher2).matchCitations(unmatched1, documents, false); verify(sparkContext).parallelizePairs(listCaptor.capture()); assertTrue(listCaptor.getValue().isEmpty()); verify(emptyRDD).union(matched1); verify(joinedMatched1).union(matched2); verify(documentAttacher).attachDocuments(joinedMatched2, documents); verify(citationAttacher).attachCitationsAndLimitDocs(documentAttached, citations); verify(bestMatchedCitationPicker).pickBest(citationAttached); }