ElasticSearch IntegrationTest

ElasticSearch IntegrationTestを使えば，メモリ上での仮想クラスタ上で単体テストを走らせることができる優れものだが，あまりサンプルコードが見つからない気がする．

package org.tanuneko;

import org.apache.commons.io.FileUtils;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.index.query.FilterBuilders;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.test.ElasticsearchIntegrationTest;
import org.elasticsearch.test.ElasticsearchIntegrationTest.ClusterScope;
import org.junit.Test;

import java.io.File;
import java.text.SimpleDateFormat;
import java.time.LocalDateTime;
import java.util.Date;
import java.util.Map;
import java.util.function.Function;

import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.hamcrest.core.Is.is;


@ClusterScope(scope= ElasticsearchIntegrationTest.Scope.SUITE,numDataNodes=2)
public class ESIntegrationTest extends ElasticsearchIntegrationTest {

    private static final String INDEX_NAME = "tanuapp";
    private static final String DOCTYPE_NAME = "posting_msg";
    private static final String TODAY_DATE_STR = new SimpleDateFormat("yyyy-MM-dd").format(new Date());

    @Test
    public void testSearch() throws Exception {
        createIndex(INDEX_NAME);
        client().admin().indices().preparePutMapping(INDEX_NAME)
                .setType(DOCTYPE_NAME)
                .setSource(readMappingFile(DOCTYPE_NAME + "_mappings.json"))
                .execute().actionGet();

        assertThat(indexDoc(INDEX_NAME, DOCTYPE_NAME, "1", "neko32", "hellou").isCreated(), is(true));
        assertThat(indexDoc(INDEX_NAME, DOCTYPE_NAME, "2", "ushi", "poooo").isCreated(), is(true));
        assertThat(indexDoc(INDEX_NAME, DOCTYPE_NAME, "3", "neko64", "ka-kakinkin ka-kinkin").isCreated(), is(true));

        flushAndRefresh(INDEX_NAME);
        ensureGreen(INDEX_NAME);
        assertThat(indexExists(INDEX_NAME), is(true));

        GetResponse resp = client().prepareGet(INDEX_NAME, DOCTYPE_NAME, "3")
                .execute()
                .actionGet();

        Map source = resp.getSource();
        assertThat((String)source.get("user"), is("neko64"));
        assertThat((String)source.get("message"), is("ka-kakinkin ka-kinkin"));
        final String postedDate = (String)source.get("postDate");
        Function isSameDate = x -> x.startsWith(TODAY_DATE_STR);
        assertTrue(isSameDate.apply(convertToGMT(postedDate)).booleanValue());

        SearchResponse searchResp = client().prepareSearch(INDEX_NAME)
                .setTypes(DOCTYPE_NAME)
                .setSearchType(SearchType.DFS_QUERY_THEN_FETCH)
                .setQuery(QueryBuilders.fuzzyLikeThisFieldQuery("user").likeText("neko50"))
                .setPostFilter(FilterBuilders.idsFilter(DOCTYPE_NAME).ids("1"))
                .execute()
                .actionGet();

        assertThat(searchResp.getHits().getHits().length, is(1));
        assertThat(searchResp.getHits().getAt(0).getSource().get("user"), is("neko32"));
    }


    private String readMappingFile(final String fileName) throws Exception {
        return FileUtils.readFileToString(new File(ESIntegrationTest.class.getClassLoader().getResource(fileName).getPath()));
    }

    private String convertToGMT(String dateStr) {
        return LocalDateTime.parse(dateStr.substring(0,dateStr.length()-1)).toString();
    }

    private IndexResponse indexDoc(final String idxName,
                                    final String docName,
                                    final String id,
                                    final String user,
                                    final String msg) throws Exception {
        return client().prepareIndex(INDEX_NAME, DOCTYPE_NAME)
                .setSource(genMsg(user, msg))
                .setId(id)
                .execute().actionGet();
    }

    private XContentBuilder genMsg(String user, String msg) throws Exception {
        return jsonBuilder().startObject()
                .field("user", user)
                .field("postDate", new Date())
                .field("message", msg)
                .endObject();
    }
}