QueryComponentTest.java

package pl.allegro.tech.search.elasticsearch.tools.reindex.process;

import com.google.common.collect.ImmutableMap;
import org.elasticsearch.action.admin.indices.settings.get.GetSettingsRequest;
import org.elasticsearch.action.admin.indices.settings.get.GetSettingsResponse;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.unit.TimeValue;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import pl.allegro.tech.search.elasticsearch.tools.reindex.connection.ElasticDataPointer;
import pl.allegro.tech.search.elasticsearch.tools.reindex.connection.ElasticSearchClientFactory;
import pl.allegro.tech.search.elasticsearch.tools.reindex.connection.ElasticSearchQuery;
import pl.allegro.tech.search.elasticsearch.tools.reindex.embeded.EmbeddedElasticsearchCluster;
import pl.allegro.tech.search.elasticsearch.tools.reindex.embeded.IndexDocument;

import java.util.stream.IntStream;
import java.util.stream.Stream;

import static org.junit.Assert.*;

public class QueryComponentTest {
  private static final String SOURCE_INDEX = "sourceindex";
  private static final String TARGET_INDEX = "targetindex";
  public static final String DATA_TYPE = "type";

  private static EmbeddedElasticsearchCluster embeddedElasticsearchCluster;

  @BeforeClass
  public static void setUp() throws Exception {
    embeddedElasticsearchCluster = EmbeddedElasticsearchCluster.createDataNode();
  }


  @AfterClass
  public static void tearDown() throws Exception {
    embeddedElasticsearchCluster.close();
  }

  @Before
  public void clearTargetIndex() {
    embeddedElasticsearchCluster.deleteIndex(SOURCE_INDEX);
    embeddedElasticsearchCluster.deleteIndex(TARGET_INDEX);
  }

  @Test
  public void testQueryNoData() {
    // given
    embeddedElasticsearchCluster.recreateIndex(SOURCE_INDEX);
    ElasticDataPointer sourceDataPointer = embeddedElasticsearchCluster.createDataPointer(SOURCE_INDEX);
    Client sourceClient = ElasticSearchClientFactory.createClient(sourceDataPointer);
    ElasticSearchQuery elasticSearchQuery = embeddedElasticsearchCluster.createInitialQuery("");

    // when
    QueryComponent component = QueryComponentBuilder.builder()
            .setClient(sourceClient)
            .setDataPointer(sourceDataPointer)
            .setQuery(elasticSearchQuery)
            .createQueryComponent();
    SearchResponse searchResponse = component.prepareSearchScrollRequest();

    // then
    assertEquals("No results overall",
            0L, searchResponse.getHits().getTotalHits());
    assertEquals("Initially zero documents are loaded",
            0L, searchResponse.getHits().getHits().length);
    assertEquals("Initially zero documents are loaded",
            0L, component.getResponseSize(searchResponse));
    assertFalse("Some documents are found",
            component.searchResultsNotEmpty(searchResponse));

    // when
    searchResponse = component.getNextScrolledSearchResults(searchResponse.getScrollId());

    // then
    assertEquals("No results overall",
            0L, searchResponse.getHits().getTotalHits());
    assertEquals("Initially zero documents are loaded",
            0L, searchResponse.getHits().getHits().length);
    assertEquals("Initially zero documents are loaded",
            0L, component.getResponseSize(searchResponse));
    assertFalse("Some documents are found",
            component.searchResultsNotEmpty(searchResponse));
  }

  @Test
  public void testQueryWithData() {
    // given
    indexWithSampleData(7000);
    ElasticDataPointer sourceDataPointer = embeddedElasticsearchCluster.createDataPointer(SOURCE_INDEX);
    Client sourceClient = ElasticSearchClientFactory.createClient(sourceDataPointer);
    ElasticSearchQuery elasticSearchQuery = embeddedElasticsearchCluster.createInitialQuery("");

    GetSettingsResponse indexSettings = sourceClient.admin().indices().getSettings(new GetSettingsRequest().indices(SOURCE_INDEX)).actionGet();
    assertEquals("We should have an index with 5 shards now",
            "5", indexSettings.getIndexToSettings().get(SOURCE_INDEX).get("index.number_of_shards"));
    assertEquals("We should have an index with one replica now",
            "1", indexSettings.getIndexToSettings().get(SOURCE_INDEX).get("index.number_of_replicas"));

    // when
    QueryComponent component = QueryComponentBuilder.builder()
            .setClient(sourceClient)
            .setDataPointer(sourceDataPointer)
            .setQuery(elasticSearchQuery)
            .createQueryComponent();
    SearchResponse searchResponse = component.prepareSearchScrollRequest();

    // then
    assertEquals("Overall there should be 7000 hits",
            7000L, searchResponse.getHits().getTotalHits());
    assertEquals("Initially zero documents are loaded",
            0L, searchResponse.getHits().getHits().length);
    assertEquals("Initially zero documents are loaded",
            0L, component.getResponseSize(searchResponse));
    assertTrue("Some documents are found",
            component.searchResultsNotEmpty(searchResponse));

    // when
    searchResponse = component.getNextScrolledSearchResults(searchResponse.getScrollId());

    // then
    assertEquals("Overall there should be 7000 hits",
            7000L, searchResponse.getHits().getTotalHits());
    assertEquals("QueryComponent tries to compute the hits to be 5000 on evenly distributed documents, never more!",
            5000L, searchResponse.getHits().getHits().length);
    assertEquals("QueryComponent tries to compute the hits to be 5000 on evenly distributed documents, never more!",
            5000L, component.getResponseSize(searchResponse));
    assertTrue("Some documents are found",
            component.searchResultsNotEmpty(searchResponse));
  }

  // just a simple test to verify that replica-shards are not included in the calculation of results
  // for the "size per shard" setting in scan/scroll-queries
  @Test
  public void testElasticsearchReplicaHandlingInScrolls() {
    // given
    indexWithSampleData(200);
    ElasticDataPointer sourceDataPointer = embeddedElasticsearchCluster.createDataPointer(SOURCE_INDEX);
    Client sourceClient = ElasticSearchClientFactory.createClient(sourceDataPointer);

    GetSettingsResponse indexSettings = sourceClient.admin().indices().getSettings(new GetSettingsRequest().indices(SOURCE_INDEX)).actionGet();
    assertEquals("We should have an index with 5 shards now",
            "5", indexSettings.getIndexToSettings().get(SOURCE_INDEX).get("index.number_of_shards"));
    assertEquals("We should have an index with one replica now",
            "1", indexSettings.getIndexToSettings().get(SOURCE_INDEX).get("index.number_of_replicas"));

    // when
    SearchRequestBuilder searchRequestBuilder = sourceClient.prepareSearch(sourceDataPointer.getIndexName())
            .setTypes(DATA_TYPE)
            .setSearchType(SearchType.SCAN)
            .addFields("_ttl", "_source")
            .setScroll(new TimeValue(QueryComponent.SCROLL_TIME_LIMIT))
            .setSize(10);
    assertNotNull(searchRequestBuilder);

    // then
    SearchResponse searchResponse = searchRequestBuilder.execute().actionGet();
    assertEquals("Overall there should be 200 hits",
            200L, searchResponse.getHits().getTotalHits());
    assertEquals("Initially zero documents are loaded",
            0L, searchResponse.getHits().getHits().length);

    // when
    searchResponse = sourceClient.prepareSearchScroll(searchResponse.getScrollId())
            .setScroll(new TimeValue(QueryComponent.SCROLL_TIMEOUT))
            .get();

    // then
    assertEquals(200L, searchResponse.getHits().getTotalHits());
    assertEquals(50L, searchResponse.getHits().getHits().length);
  }


  private void indexWithSampleData(final int numberOfDocuments) {
    Stream<IndexDocument> streamToBeIndexed = IntStream
            .range(1, numberOfDocuments+1)
            .mapToObj(
                    i -> new IndexDocument(Integer.toString(i), ImmutableMap.of("fieldName", i))
            );
    embeddedElasticsearchCluster.indexWithSampleData(SOURCE_INDEX, DATA_TYPE, streamToBeIndexed);
  }
}