Examples with FetchPhase org.opensearch.search.fetch.FetchPhase used on opensource projects

Example 1 with FetchPhase

use of org.opensearch.search.fetch.FetchPhase in project OpenSearch by opensearch-project.

the class AggregatorTestCase method createSearchContext.

protected SearchContext createSearchContext(IndexSearcher indexSearcher, IndexSettings indexSettings, Query query, MultiBucketConsumer bucketConsumer, CircuitBreakerService circuitBreakerService, MappedFieldType... fieldTypes) throws IOException {
    QueryCache queryCache = new DisabledQueryCache(indexSettings);
    QueryCachingPolicy queryCachingPolicy = new QueryCachingPolicy() {

        @Override
        public void onUse(Query query) {
        }

        @Override
        public boolean shouldCache(Query query) {
            // never cache a query
            return false;
        }
    };
    ContextIndexSearcher contextIndexSearcher = new ContextIndexSearcher(indexSearcher.getIndexReader(), indexSearcher.getSimilarity(), queryCache, queryCachingPolicy, false);
    SearchContext searchContext = mock(SearchContext.class);
    when(searchContext.numberOfShards()).thenReturn(1);
    when(searchContext.searcher()).thenReturn(contextIndexSearcher);
    when(searchContext.fetchPhase()).thenReturn(new FetchPhase(Arrays.asList(new FetchSourcePhase(), new FetchDocValuesPhase())));
    when(searchContext.bitsetFilterCache()).thenReturn(new BitsetFilterCache(indexSettings, mock(Listener.class)));
    IndexShard indexShard = mock(IndexShard.class);
    when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0));
    when(searchContext.indexShard()).thenReturn(indexShard);
    when(searchContext.aggregations()).thenReturn(new SearchContextAggregations(AggregatorFactories.EMPTY, bucketConsumer));
    when(searchContext.query()).thenReturn(query);
    /*
         * Always use the circuit breaking big arrays instance so that the CircuitBreakerService
         * we're passed gets a chance to break.
         */
    BigArrays bigArrays = new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), circuitBreakerService).withCircuitBreaking();
    when(searchContext.bigArrays()).thenReturn(bigArrays);
    // TODO: now just needed for top_hits, this will need to be revised for other agg unit tests:
    MapperService mapperService = mapperServiceMock();
    when(mapperService.getIndexSettings()).thenReturn(indexSettings);
    when(mapperService.hasNested()).thenReturn(false);
    when(searchContext.mapperService()).thenReturn(mapperService);
    IndexFieldDataService ifds = new IndexFieldDataService(indexSettings, new IndicesFieldDataCache(Settings.EMPTY, new IndexFieldDataCache.Listener() {
    }), circuitBreakerService, mapperService);
    QueryShardContext queryShardContext = queryShardContextMock(contextIndexSearcher, mapperService, indexSettings, circuitBreakerService, bigArrays);
    when(searchContext.getQueryShardContext()).thenReturn(queryShardContext);
    when(queryShardContext.getObjectMapper(anyString())).thenAnswer(invocation -> {
        String fieldName = (String) invocation.getArguments()[0];
        if (fieldName.startsWith(NESTEDFIELD_PREFIX)) {
            BuilderContext context = new BuilderContext(indexSettings.getSettings(), new ContentPath());
            return new ObjectMapper.Builder<>(fieldName).nested(Nested.newNested()).build(context);
        }
        return null;
    });
    Map<String, MappedFieldType> fieldNameToType = new HashMap<>();
    fieldNameToType.putAll(Arrays.stream(fieldTypes).filter(Objects::nonNull).collect(Collectors.toMap(MappedFieldType::name, Function.identity())));
    fieldNameToType.putAll(getFieldAliases(fieldTypes));
    registerFieldTypes(searchContext, mapperService, fieldNameToType);
    doAnswer(invocation -> {
        /* Store the release-ables so we can release them at the end of the test case. This is important because aggregations don't
             * close their sub-aggregations. This is fairly similar to what the production code does. */
        releasables.add((Releasable) invocation.getArguments()[0]);
        return null;
    }).when(searchContext).addReleasable(any());
    return searchContext;
}