/solr/core/src/test/org/apache/solr/handler/component/CustomHighlightComponentTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.solr.handler.component;

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.request.CollectionAdminRequest;
import org.apache.solr.client.solrj.request.QueryRequest;
import org.apache.solr.client.solrj.request.UpdateRequest;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.cloud.AbstractDistribZkTestBase;
import org.apache.solr.cloud.ConfigRequest;
import org.apache.solr.cloud.SolrCloudTestCase;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.common.util.SimpleOrderedMap;
import org.apache.solr.highlight.SolrFragmentsBuilder;
import org.junit.BeforeClass;
import org.junit.Test;

public class CustomHighlightComponentTest extends SolrCloudTestCase {

  public static class CustomHighlightComponent extends HighlightComponent {

    protected String id_key = "id";
    protected String snippets_key = "snippets";

    @Override
    protected String highlightingResponseField() {
      return "custom_highlighting";
    }

    @Override
    protected Object convertHighlights(NamedList<Object> hl) {
      final ArrayList<SimpleOrderedMap<Object>> hlMaps = new ArrayList<>();
      for (int i = 0; i < hl.size(); ++i) {
        SimpleOrderedMap<Object> hlMap = new SimpleOrderedMap<>();
        hlMap.add(id_key, hl.getName(i));
        hlMap.add(snippets_key, hl.getVal(i));
        hlMaps.add(hlMap);
      }
      return hlMaps;
    }

    @Override
    protected Object[] newHighlightsArray(int size) {
      return (SimpleOrderedMap<?>[]) Array.newInstance(SimpleOrderedMap.class, size);
    }

    @Override
    protected void addHighlights(Object[] objArr, Object obj, Map<Object, ShardDoc> resultIds) {
      SimpleOrderedMap<?>[] mapArr = (SimpleOrderedMap<?>[]) objArr;
      @SuppressWarnings("unchecked")
      final ArrayList<SimpleOrderedMap<?>> hlMaps = (ArrayList<SimpleOrderedMap<?>>) obj;
      for (SimpleOrderedMap<?> hlMap : hlMaps) {
        String id = (String) hlMap.get(id_key);
        ShardDoc sdoc = resultIds.get(id);
        int idx = sdoc.positionInResponse;
        mapArr[idx] = hlMap;
      }
    }

    @Override
    protected Object getAllHighlights(Object[] objArr) {
      final SimpleOrderedMap<?>[] mapArr = (SimpleOrderedMap<?>[]) objArr;
      // remove nulls in case not all docs were able to be retrieved
      ArrayList<SimpleOrderedMap<?>> mapList = new ArrayList<>();
      for (SimpleOrderedMap<?> map : mapArr) {
        if (map != null) {
          mapList.add(map);
        }
      }
      return mapList;
    }
  }

  protected String customHighlightComponentClassName() {
    return CustomHighlightComponent.class.getName();
  }

  protected String id_key = "id";
  protected String snippets_key = "snippets";

  private static String COLLECTION;

  @BeforeClass
  public static void setupCluster() throws Exception {

    // decide collection name ...
    COLLECTION = "collection" + (1 + random().nextInt(100));
    // ... and shard/replica/node numbers
    final int numShards = 3;
    final int numReplicas = 2;
    final int nodeCount = numShards * numReplicas;

    // create and configure cluster
    configureCluster(nodeCount).addConfig("conf", configset("cloud-dynamic")).configure();

    // create an empty collection
    CollectionAdminRequest.createCollection(COLLECTION, "conf", numShards, numReplicas)
        .processAndWait(cluster.getSolrClient(), DEFAULT_TIMEOUT);
    AbstractDistribZkTestBase.waitForRecoveriesToFinish(
        COLLECTION, cluster.getZkStateReader(), false, true, DEFAULT_TIMEOUT);
  }

  @Test
  public void test() throws Exception {

    // determine custom search handler name (the exact name should not matter)
    final String customSearchHandlerName = "/custom_select" + random().nextInt();

    final String defaultHighlightComponentName = HighlightComponent.COMPONENT_NAME;
    final String highlightComponentName;

    // add custom component (if needed) and handler
    {
      if (random().nextBoolean()) {
        // default component
        highlightComponentName = defaultHighlightComponentName;
      } else {
        // custom component
        highlightComponentName = "customhighlight" + random().nextInt();
        cluster
            .getSolrClient()
            .request(
                new ConfigRequest(
                    "{\n"
                        + "  'add-searchcomponent': {\n"
                        + "    'name': '"
                        + highlightComponentName
                        + "',\n"
                        + "    'class': '"
                        + customHighlightComponentClassName()
                        + "'\n"
                        + "  }\n"
                        + "}"),
                COLLECTION);
      }
      // handler
      cluster
          .getSolrClient()
          .request(
              new ConfigRequest(
                  "{\n"
                      + "  'add-requesthandler': {\n"
                      + "    'name' : '"
                      + customSearchHandlerName
                      + "',\n"
                      + "    'class' : 'org.apache.solr.handler.component.SearchHandler',\n"
                      + "    'components' : [ '"
                      + QueryComponent.COMPONENT_NAME
                      + "', '"
                      + highlightComponentName
                      + "' ]\n"
                      + "  }\n"
                      + "}"),
              COLLECTION);
    }

    // add some documents
    final String id = "id";
    final String t1 = "a_t";
    final String t2 = "b_t";
    {
      new UpdateRequest()
          .add(sdoc(id, 1, t1, "bumble bee", t2, "bumble bee"))
          .add(sdoc(id, 2, t1, "honey bee", t2, "honey bee"))
          .add(sdoc(id, 3, t1, "solitary bee", t2, "solitary bee"))
          .commit(cluster.getSolrClient(), COLLECTION);
    }

    // search for the documents
    {
      // compose the query
      final SolrQuery solrQuery = new SolrQuery(t1 + ":bee");
      solrQuery.setRequestHandler(customSearchHandlerName);
      solrQuery.setHighlight(true);
      final boolean t1Highlights = random().nextBoolean();
      if (t1Highlights) {
        solrQuery.addHighlightField(t1);
      }
      final boolean t2Highlights = random().nextBoolean();
      if (t2Highlights) {
        solrQuery.addHighlightField(t2);
      }

      // make the query
      final QueryResponse queryResponse =
          new QueryRequest(solrQuery).process(cluster.getSolrClient(), COLLECTION);

      // analyse the response
      final Map<String, Map<String, List<String>>> highlighting = queryResponse.getHighlighting();
      @SuppressWarnings({"unchecked"})
      final ArrayList<SimpleOrderedMap<Object>> custom_highlighting =
          (ArrayList<SimpleOrderedMap<Object>>)
              queryResponse.getResponse().get("custom_highlighting");

      if (defaultHighlightComponentName.equals(highlightComponentName)) {
        // regular 'highlighting' ...
        if (t1Highlights) {
          checkHighlightingResponseMap(highlighting, t1);
        }
        if (t2Highlights) {
          checkHighlightingResponseMap(highlighting, t2);
        }
        if (!t1Highlights && !t2Highlights) {
          checkHighlightingResponseMap(highlighting, null);
        }
        // ... and no 'custom_highlighting'
        assertNull(custom_highlighting);
      } else {
        // no regular 'highlighting' ...
        assertNull(highlighting);
        // ... but 'custom_highlighting'
        assertNotNull(custom_highlighting);
        if (t1Highlights) {
          checkHighlightingResponseList(custom_highlighting, t1);
        }
        if (t2Highlights) {
          checkHighlightingResponseList(custom_highlighting, t2);
        }
        if (!t1Highlights && !t2Highlights) {
          checkHighlightingResponseList(custom_highlighting, null);
        }
      }
    }
  }

  protected void checkHighlightingResponseMap(
      Map<String, Map<String, List<String>>> highlightingMap, String highlightedField) {
    assertEquals(
        "too few or too many keys: " + highlightingMap.keySet(), 3, highlightingMap.size());
    checkHighlightingResponseMapElement(
        highlightingMap.get("1"), highlightedField, "bumble ", "bee");
    checkHighlightingResponseMapElement(
        highlightingMap.get("2"), highlightedField, "honey ", "bee");
    checkHighlightingResponseMapElement(
        highlightingMap.get("3"), highlightedField, "solitary ", "bee");
  }

  protected void checkHighlightingResponseMapElement(
      Map<String, List<String>> docHighlights,
      String highlightedField,
      String preHighlightText,
      String highlightedText) {
    if (highlightedField == null) {
      assertEquals(0, docHighlights.size());
    } else {
      List<String> docHighlightsList = docHighlights.get(highlightedField);
      assertEquals(1, docHighlightsList.size());
      assertEquals(
          preHighlightText
              + SolrFragmentsBuilder.DEFAULT_PRE_TAGS
              + highlightedText
              + SolrFragmentsBuilder.DEFAULT_POST_TAGS,
          docHighlightsList.get(0));
    }
  }

  protected void checkHighlightingResponseList(
      ArrayList<SimpleOrderedMap<Object>> highlightingList, String highlightedField) {
    assertEquals(
        "too few or too many elements: " + highlightingList.size(), 3, highlightingList.size());
    final Set<String> seenDocIds = new HashSet<>();
    for (SimpleOrderedMap<Object> highlightingListElementMap : highlightingList) {
      final String expectedHighlightText;
      final String actualHighlightText;
      // two elements in total: id and snippets
      assertEquals(highlightingList.toString(), 2, highlightingListElementMap.size());
      // id element
      {
        final String docId = (String) highlightingListElementMap.get(id_key);
        seenDocIds.add(docId);
        final String preHighlightText;
        final String highlightedText = "bee";
        if ("1".equals(docId)) {
          preHighlightText = "bumble ";
        } else if ("2".equals(docId)) {
          preHighlightText = "honey ";
        } else if ("3".equals(docId)) {
          preHighlightText = "solitary ";
        } else {
          preHighlightText = null;
          fail("unknown docId " + docId);
        }
        expectedHighlightText =
            preHighlightText
                + SolrFragmentsBuilder.DEFAULT_PRE_TAGS
                + highlightedText
                + SolrFragmentsBuilder.DEFAULT_POST_TAGS;
      }
      // snippets element
      {
        @SuppressWarnings({"unchecked"})
        SimpleOrderedMap<Object> snippets =
            (SimpleOrderedMap<Object>) highlightingListElementMap.get(snippets_key);
        if (highlightedField == null) {
          assertEquals(0, snippets.size());
        } else {
          @SuppressWarnings({"unchecked"})
          ArrayList<String> docHighlights = (ArrayList<String>) (snippets).get(highlightedField);
          assertEquals(1, docHighlights.size());
          actualHighlightText = docHighlights.get(0);
          assertEquals(expectedHighlightText, actualHighlightText);
        }
      }
    }
    assertEquals(3, seenDocIds.size());
  }
}