@Test

public void shouldReturnHighlightedFieldsForGivenQueryAndFields() {

    //given

    String documentId = randomNumeric(5);

    String actualMessage = "some test message";

    String highlightedMessage = "some <em>test</em> message";

    SampleEntity sampleEntity = SampleEntity.builder().id(documentId)

            .message(actualMessage)

            .version(System.currentTimeMillis()).build();

    IndexQuery indexQuery = getIndexQuery(sampleEntity);

    elasticsearchTemplate.index(indexQuery);

    elasticsearchTemplate.refresh(SampleEntity.class);

    SearchQuery searchQuery = new NativeSearchQueryBuilder()

            .withQuery(termQuery("message", "test"))

            .withHighlightFields(new HighlightBuilder.Field("message"))

            .build();

    Page<SampleEntity> sampleEntities = elasticsearchTemplate.queryForPage(searchQuery, SampleEntity.class, new SearchResultMapper() {

        @Override

        public <T> Page<T> mapResults(SearchResponse response, Class<T> clazz, Pageable pageable) {

            List<SampleEntity> chunk = new ArrayList<SampleEntity>();

            for (SearchHit searchHit : response.getHits()) {

                if (response.getHits().getHits().length <= 0) {

                    return null;

                }

                SampleEntity user = new SampleEntity();

                user.setId(searchHit.getId());

                user.setMessage((String) searchHit.getSource().get("message"));

                user.setHighlightedMessage(searchHit.getHighlightFields().get("message").fragments()[0].toString());

                chunk.add(user);

            }

            if (chunk.size() > 0) {

                return new PageImpl<T>((List<T>) chunk);

            }

            return null;

        }

    });

    assertThat(sampleEntities.getContent().get(0).getHighlightedMessage(), is(highlightedMessage));

}