How to provide highlighting with Spring data elasticsearch

How to provide highlighting with Spring data elasticsearch

@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));
}