@Override public void query(Query query, int limit, IndexResultCollector collector) throws IOException { long start = stats.startRepositoryQuery(); int totalHits = 0; IndexSearcher searcher = searcherManager.acquire(); try { TopDocs docs = searcher.search(query, limit); totalHits = docs.totalHits; for (ScoreDoc scoreDoc : docs.scoreDocs) { Document doc = searcher.doc(scoreDoc.doc); Object key = SerializerUtil.getKey(doc); if (logger.isDebugEnabled()) { logger.debug("query found doc:" + doc + ":" + scoreDoc); } collector.collect(key, scoreDoc.score); } } finally { searcherManager.release(searcher); stats.endRepositoryQuery(start, totalHits); } }
@Test public void injectCustomCollectorManager() throws Exception { final CollectorManager mockManager = mock(CollectorManager.class); searchArgs = new LuceneFunctionContext<IndexResultCollector>(queryProvider, "indexName", mockManager); when(mockContext.getDataSet()).thenReturn(mockRegion); when(mockContext.getArguments()).thenReturn(searchArgs); when(mockContext.getResultSender()).thenReturn(mockResultSender); repos.remove(0); when(mockRepoManager.getRepositories(eq(mockContext), eq(false))).thenReturn(repos); when(mockManager.newCollector(eq("repo2"))).thenReturn(mockCollector); when(mockManager.reduce(any(Collection.class))).thenAnswer(invocation -> { Collection<IndexResultCollector> collectors = invocation.getArgument(0); assertEquals(1, collectors.size()); assertEquals(mockCollector, collectors.iterator().next()); return new TopEntriesCollector(null); }); doAnswer(invocation -> { IndexResultCollector collector = invocation.getArgument(2); collector.collect(r2_1.getKey(), r2_1.getScore()); return null; }).when(mockRepository2).query(eq(query), eq(LuceneQueryFactory.DEFAULT_LIMIT), any(IndexResultCollector.class)); LuceneQueryFunction function = new LuceneQueryFunction(); function.execute(mockContext); verify(mockCollector).collect(eq("key-2-1"), eq(.45f)); verify(mockResultSender).lastResult(any(TopEntriesCollector.class)); }
doAnswer(invocation -> { IndexResultCollector collector = invocation.getArgument(2); collector.collect(r1_1.getKey(), r1_1.getScore()); collector.collect(r1_2.getKey(), r1_2.getScore()); collector.collect(r1_3.getKey(), r1_3.getScore()); return null; }).when(mockRepository1).query(eq(query), eq(LuceneQueryFactory.DEFAULT_LIMIT), collector.collect(r2_1.getKey(), r2_1.getScore()); collector.collect(r2_2.getKey(), r2_2.getScore()); return null; }).when(mockRepository2).query(eq(query), eq(LuceneQueryFactory.DEFAULT_LIMIT),
collector.collect(r1_1.getKey(), r1_1.getScore()); collector.collect(r1_2.getKey(), r1_2.getScore()); collector.collect(r1_3.getKey(), r1_3.getScore()); return null; }).when(mockRepository1).query(eq(query), eq(3), any(IndexResultCollector.class)); collector.collect(r2_1.getKey(), r2_1.getScore()); collector.collect(r2_2.getKey(), r2_2.getScore()); return null; }).when(mockRepository2).query(eq(query), eq(3), any(IndexResultCollector.class));