org.apache.metron.solr.dao.SolrSearchDaoTest.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.metron.solr.dao.SolrSearchDaoTest.java

Source

/**
 * 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.metron.solr.dao;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.core.IsCollectionContaining.hasItems;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.argThat;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
import static org.powermock.api.mockito.PowerMockito.mockStatic;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.apache.metron.common.Constants;
import org.apache.metron.indexing.dao.AccessConfig;
import org.apache.metron.indexing.dao.search.GetRequest;
import org.apache.metron.indexing.dao.search.Group;
import org.apache.metron.indexing.dao.search.GroupOrder;
import org.apache.metron.indexing.dao.search.GroupRequest;
import org.apache.metron.indexing.dao.search.GroupResponse;
import org.apache.metron.indexing.dao.search.GroupResult;
import org.apache.metron.indexing.dao.search.InvalidSearchException;
import org.apache.metron.indexing.dao.search.SearchRequest;
import org.apache.metron.indexing.dao.search.SearchResponse;
import org.apache.metron.indexing.dao.search.SearchResult;
import org.apache.metron.indexing.dao.search.SortField;
import org.apache.metron.indexing.dao.update.Document;
import org.apache.metron.solr.matcher.ModifiableSolrParamsMatcher;
import org.apache.metron.solr.matcher.SolrQueryMatcher;
import org.apache.solr.client.solrj.SolrClient;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.request.CollectionAdminRequest;
import org.apache.solr.client.solrj.response.FacetField;
import org.apache.solr.client.solrj.response.FieldStatsInfo;
import org.apache.solr.client.solrj.response.PivotField;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.common.util.NamedList;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;

@RunWith(PowerMockRunner.class)
@PrepareForTest({ CollectionAdminRequest.class })
public class SolrSearchDaoTest {

    @Rule
    public final ExpectedException exception = ExpectedException.none();

    private SolrClient client;
    private AccessConfig accessConfig;
    private SolrSearchDao solrSearchDao;
    private SolrRetrieveLatestDao solrRetrieveLatestDao;

    @SuppressWarnings("unchecked")
    @Before
    public void setUp() throws Exception {
        client = mock(SolrClient.class);
        accessConfig = mock(AccessConfig.class);
        when(accessConfig.getIndexSupplier()).thenReturn(sensorType -> sensorType);
        solrSearchDao = new SolrSearchDao(client, accessConfig);
        solrRetrieveLatestDao = new SolrRetrieveLatestDao(client, accessConfig);
        mockStatic(CollectionAdminRequest.class);
        when(CollectionAdminRequest.listCollections(client)).thenReturn(Arrays.asList("bro", "snort"));
    }

    @Test
    public void searchShouldProperlyReturnSearchResponse() throws Exception {
        SearchRequest searchRequest = mock(SearchRequest.class);
        SearchResponse searchResponse = mock(SearchResponse.class);
        SolrQuery solrQuery = mock(SolrQuery.class);
        QueryResponse queryResponse = mock(QueryResponse.class);

        solrSearchDao = spy(new SolrSearchDao(client, accessConfig));
        when(searchRequest.getQuery()).thenReturn("query");
        doReturn(solrQuery).when(solrSearchDao).buildSearchRequest(searchRequest, "*");
        when(client.query(solrQuery)).thenReturn(queryResponse);
        doReturn(searchResponse).when(solrSearchDao).buildSearchResponse(searchRequest, queryResponse);

        assertEquals(searchResponse, solrSearchDao.search(searchRequest, "*"));
        verify(solrSearchDao).buildSearchRequest(searchRequest, "*");
        verify(client).query(solrQuery);
        verify(solrSearchDao).buildSearchResponse(searchRequest, queryResponse);
        verifyNoMoreInteractions(client);
    }

    @Test
    public void searchShouldThrowInvalidSearchExceptionOnEmptyQuery() throws Exception {
        exception.expect(InvalidSearchException.class);
        exception.expectMessage("Search query is invalid: null");

        solrSearchDao.search(new SearchRequest());
    }

    @Test
    public void searchShouldThrowInvalidSearchExceptionOnEmptyClient() throws Exception {
        exception.expect(InvalidSearchException.class);
        exception.expectMessage("Uninitialized Dao!  You must call init() prior to use.");

        SearchRequest searchRequest = new SearchRequest();
        searchRequest.setQuery("query");
        new SolrSearchDao(null, accessConfig).search(searchRequest);
    }

    @Test
    public void searchShouldThrowSearchResultSizeException() throws Exception {
        exception.expect(InvalidSearchException.class);
        exception.expectMessage("Search result size must be less than 100");

        when(accessConfig.getMaxSearchResults()).thenReturn(100);
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.setQuery("query");
        searchRequest.setSize(200);
        solrSearchDao.search(searchRequest);
    }

    @Test
    public void groupShouldProperlyReturnGroupResponse() throws Exception {
        GroupRequest groupRequest = mock(GroupRequest.class);
        QueryResponse queryResponse = mock(QueryResponse.class);
        GroupResponse groupResponse = mock(GroupResponse.class);

        solrSearchDao = spy(new SolrSearchDao(client, accessConfig));
        Group group1 = new Group();
        group1.setField("field1");
        Group group2 = new Group();
        group2.setField("field2");
        when(groupRequest.getQuery()).thenReturn("query");
        when(groupRequest.getGroups()).thenReturn(Arrays.asList(group1, group2));
        when(groupRequest.getScoreField()).thenReturn(Optional.of("scoreField"));
        when(groupRequest.getIndices()).thenReturn(Arrays.asList("bro", "snort"));
        when(client.query(any())).thenReturn(queryResponse);
        doReturn(groupResponse).when(solrSearchDao).buildGroupResponse(groupRequest, queryResponse);
        SolrQuery expectedSolrQuery = new SolrQuery().setStart(0).setRows(0).setQuery("query");
        expectedSolrQuery.set("collection", "bro,snort");
        expectedSolrQuery.set("stats", true);
        expectedSolrQuery.set("stats.field", "{!tag=piv1 sum=true}scoreField");
        expectedSolrQuery.set("facet", true);
        expectedSolrQuery.set("facet.pivot", "{!stats=piv1}field1,field2");

        assertEquals(groupResponse, solrSearchDao.group(groupRequest));
        verify(client).query(argThat(new SolrQueryMatcher(expectedSolrQuery)));
        verify(solrSearchDao).buildGroupResponse(groupRequest, queryResponse);

        verifyNoMoreInteractions(client);
    }

    @Test
    public void getLatestShouldProperlyReturnDocument() throws Exception {
        SolrDocument solrDocument = mock(SolrDocument.class);

        solrSearchDao = spy(new SolrSearchDao(client, accessConfig));
        when(client.getById("collection", "guid")).thenReturn(solrDocument);
        Document document = SolrUtilities.toDocument(solrDocument);

        assertEquals(document, solrRetrieveLatestDao.getLatest("guid", "collection"));

        verify(client).getById("collection", "guid");
        verifyNoMoreInteractions(client);
    }

    @Test
    public void getAllLatestShouldProperlyReturnDocuments() throws Exception {
        GetRequest broRequest1 = new GetRequest("bro-1", "bro");
        GetRequest broRequest2 = new GetRequest("bro-2", "bro");
        GetRequest snortRequest1 = new GetRequest("snort-1", "snort");
        GetRequest snortRequest2 = new GetRequest("snort-2", "snort");
        SolrDocument broSolrDoc1 = mock(SolrDocument.class);
        SolrDocument broSolrDoc2 = mock(SolrDocument.class);
        SolrDocument snortSolrDoc1 = mock(SolrDocument.class);
        SolrDocument snortSolrDoc2 = mock(SolrDocument.class);
        Document broDoc1 = SolrUtilities.toDocument(broSolrDoc1);
        Document broDoc2 = SolrUtilities.toDocument(broSolrDoc2);
        Document snortDoc1 = SolrUtilities.toDocument(snortSolrDoc1);
        Document snortDoc2 = SolrUtilities.toDocument(snortSolrDoc2);

        solrSearchDao = spy(new SolrSearchDao(client, accessConfig));
        SolrDocumentList broList = new SolrDocumentList();
        broList.add(broSolrDoc1);
        broList.add(broSolrDoc2);
        SolrDocumentList snortList = new SolrDocumentList();
        snortList.add(snortSolrDoc1);
        snortList.add(snortSolrDoc2);
        when(client.getById((Collection<String>) argThat(hasItems("bro-1", "bro-2")),
                argThat(new ModifiableSolrParamsMatcher(new ModifiableSolrParams().set("collection", "bro")))))
                        .thenReturn(broList);
        when(client.getById((Collection<String>) argThat(hasItems("snort-1", "snort-2")),
                argThat(new ModifiableSolrParamsMatcher(new ModifiableSolrParams().set("collection", "snort")))))
                        .thenReturn(snortList);
        assertEquals(Arrays.asList(broDoc1, broDoc2, snortDoc1, snortDoc2), solrRetrieveLatestDao
                .getAllLatest(Arrays.asList(broRequest1, broRequest2, snortRequest1, snortRequest2)));
    }

    @Test
    public void buildSearchRequestShouldReturnSolrQuery() throws Exception {
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.setIndices(Arrays.asList("bro", "snort"));
        searchRequest.setSize(5);
        searchRequest.setFrom(10);
        searchRequest.setQuery("query");
        SortField sortField = new SortField();
        sortField.setField("sortField");
        sortField.setSortOrder("ASC");
        searchRequest.setSort(Collections.singletonList(sortField));
        searchRequest.setFields(Arrays.asList("field1", "field2"));
        searchRequest.setFacetFields(Arrays.asList("facetField1", "facetField2"));

        SolrQuery exceptedSolrQuery = new SolrQuery().setStart(10).setRows(5).setQuery("query")
                .addSort("sortField", SolrQuery.ORDER.asc).addField("field1").addField("field2")
                .addFacetField("facetField1", "facetField2");
        exceptedSolrQuery.set("collection", "bro,snort");

        SolrQuery solrQuery = solrSearchDao.buildSearchRequest(searchRequest, "field1,field2");
        assertThat(solrQuery, new SolrQueryMatcher(exceptedSolrQuery));
    }

    @Test
    public void buildSearchResponseShouldReturnSearchResponse() {
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.setFields(Collections.singletonList("id"));
        searchRequest.setFacetFields(Collections.singletonList("facetField"));
        QueryResponse queryResponse = mock(QueryResponse.class);
        SolrDocument solrDocument1 = new SolrDocument();
        solrDocument1.setField(Constants.GUID, "id1");
        solrDocument1.setField("id", "id1");
        SolrDocument solrDocument2 = new SolrDocument();
        solrDocument2.setField(Constants.GUID, "id2");
        solrDocument2.setField("id", "id2");

        solrSearchDao = spy(new SolrSearchDao(client, accessConfig));
        SolrDocumentList solrDocumentList = new SolrDocumentList();
        solrDocumentList.add(solrDocument1);
        solrDocumentList.add(solrDocument2);
        solrDocumentList.setNumFound(100);
        when(queryResponse.getResults()).thenReturn(solrDocumentList);
        SearchResult searchResult1 = new SearchResult();
        searchResult1.setId("id1");
        HashMap<String, Object> source1 = new HashMap<>();
        source1.put("id", "id1");
        searchResult1.setSource(source1);
        SearchResult searchResult2 = new SearchResult();
        searchResult2.setId("id2");
        HashMap<String, Object> source2 = new HashMap<>();
        source2.put("id", "id2");
        searchResult2.setSource(source2);
        Map<String, Map<String, Long>> facetCounts = new HashMap<String, Map<String, Long>>() {
            {
                put("id", new HashMap<String, Long>() {
                    {
                        put("id1", 1L);
                        put("id2", 1L);
                    }
                });
            }
        };
        doReturn(facetCounts).when(solrSearchDao).getFacetCounts(Collections.singletonList("facetField"),
                queryResponse);
        SearchResponse expectedSearchResponse = new SearchResponse();
        SearchResult expectedSearchResult1 = new SearchResult();
        expectedSearchResult1.setId("id1");
        expectedSearchResult1.setSource(source1);
        SearchResult expectedSearchResult2 = new SearchResult();
        expectedSearchResult2.setId("id2");
        expectedSearchResult2.setSource(source2);

        expectedSearchResponse.setResults(Arrays.asList(expectedSearchResult1, expectedSearchResult2));
        expectedSearchResponse.setTotal(100);
        expectedSearchResponse.setFacetCounts(facetCounts);

        assertEquals(expectedSearchResponse, solrSearchDao.buildSearchResponse(searchRequest, queryResponse));
    }

    @Test
    public void getSearchResultShouldProperlyReturnResults() {
        SolrDocument solrDocument = mock(SolrDocument.class);

        when(solrDocument.getFieldValue(Constants.GUID)).thenReturn("guid");
        when(solrDocument.getFieldValue(Constants.SENSOR_TYPE)).thenReturn("sensorType");
        when(solrDocument.getFieldValue("field1")).thenReturn("value1");
        when(solrDocument.getFieldValue("field2")).thenReturn("value2");
        when(solrDocument.getFieldNames()).thenReturn(Arrays.asList("field1", "field2"));

        SearchResult expectedSearchResult = new SearchResult();
        expectedSearchResult.setId("guid");
        expectedSearchResult.setIndex("sensorType");
        expectedSearchResult.setSource(new HashMap<String, Object>() {
            {
                put("field1", "value1");
            }
        });

        assertEquals(expectedSearchResult, SolrUtilities.getSearchResult(solrDocument,
                Collections.singletonList("field1"), solrSearchDao.getAccessConfig().getIndexSupplier()));

        SearchResult expectedSearchResultAllFields = new SearchResult();
        expectedSearchResultAllFields.setId("guid");
        expectedSearchResultAllFields.setIndex("sensorType");
        expectedSearchResultAllFields.setSource(new HashMap<String, Object>() {
            {
                put("field1", "value1");
                put("field2", "value2");
            }
        });

        assertEquals(expectedSearchResultAllFields, SolrUtilities.getSearchResult(solrDocument, null,
                solrSearchDao.getAccessConfig().getIndexSupplier()));
    }

    @Test
    public void getFacetCountsShouldProperlyReturnFacetCounts() {
        QueryResponse queryResponse = mock(QueryResponse.class);

        FacetField facetField1 = new FacetField("field1");
        facetField1.add("value1", 1);
        facetField1.add("value2", 2);
        FacetField facetField2 = new FacetField("field2");
        facetField2.add("value3", 3);
        facetField2.add("value4", 4);
        when(queryResponse.getFacetField("field1")).thenReturn(facetField1);
        when(queryResponse.getFacetField("field2")).thenReturn(facetField2);

        Map<String, Map<String, Long>> expectedFacetCounts = new HashMap<String, Map<String, Long>>() {
            {
                put("field1", new HashMap<String, Long>() {
                    {
                        put("value1", 1L);
                        put("value2", 2L);
                    }
                });
                put("field2", new HashMap<String, Long>() {
                    {
                        put("value3", 3L);
                        put("value4", 4L);
                    }
                });
            }
        };

        assertEquals(expectedFacetCounts,
                solrSearchDao.getFacetCounts(Arrays.asList("field1", "field2"), queryResponse));
    }

    @Test
    public void buildGroupResponseShouldProperlyReturnGroupReponse() {
        GroupRequest groupRequest = mock(GroupRequest.class);
        QueryResponse queryResponse = mock(QueryResponse.class);
        NamedList namedList = mock(NamedList.class);
        List pivotFields = mock(List.class);
        List groupResults = mock(List.class);

        solrSearchDao = spy(new SolrSearchDao(client, accessConfig));
        Group group1 = new Group();
        group1.setField("field1");
        Group group2 = new Group();
        group2.setField("field2");
        when(groupRequest.getGroups()).thenReturn(Arrays.asList(group1, group2));
        when(queryResponse.getFacetPivot()).thenReturn(namedList);
        when(namedList.get("field1,field2")).thenReturn(pivotFields);
        doReturn(groupResults).when(solrSearchDao).getGroupResults(groupRequest, 0, pivotFields);

        GroupResponse groupResponse = solrSearchDao.buildGroupResponse(groupRequest, queryResponse);
        assertEquals("field1", groupResponse.getGroupedBy());
        verify(namedList).get("field1,field2");
        verify(solrSearchDao).getGroupResults(groupRequest, 0, pivotFields);

    }

    @Test
    public void getGroupResultsShouldProperlyReturnGroupResults() {
        GroupRequest groupRequest = new GroupRequest();
        Group group1 = new Group();
        group1.setField("field1");
        GroupOrder groupOrder1 = new GroupOrder();
        groupOrder1.setSortOrder("ASC");
        groupOrder1.setGroupOrderType("TERM");
        group1.setOrder(groupOrder1);
        Group group2 = new Group();
        group2.setField("field2");
        GroupOrder groupOrder2 = new GroupOrder();
        groupOrder2.setSortOrder("DESC");
        groupOrder2.setGroupOrderType("COUNT");
        group2.setOrder(groupOrder2);
        groupRequest.setGroups(Arrays.asList(group1, group2));
        groupRequest.setScoreField("score");

        PivotField level1Pivot1 = mock(PivotField.class);
        PivotField level1Pivot2 = mock(PivotField.class);
        PivotField level2Pivot1 = mock(PivotField.class);
        PivotField level2Pivot2 = mock(PivotField.class);
        FieldStatsInfo level1Pivot1FieldStatsInfo = mock(FieldStatsInfo.class);
        FieldStatsInfo level1Pivot2FieldStatsInfo = mock(FieldStatsInfo.class);
        FieldStatsInfo level2Pivot1FieldStatsInfo = mock(FieldStatsInfo.class);
        FieldStatsInfo level2Pivot2FieldStatsInfo = mock(FieldStatsInfo.class);
        List<PivotField> level1Pivots = Arrays.asList(level1Pivot1, level1Pivot2);
        List<PivotField> level2Pivots = Arrays.asList(level2Pivot1, level2Pivot2);

        when(level1Pivot1.getValue()).thenReturn("field1value1");
        when(level1Pivot1.getCount()).thenReturn(1);
        when(level1Pivot1FieldStatsInfo.getSum()).thenReturn(1.0);
        when(level1Pivot1.getFieldStatsInfo()).thenReturn(new HashMap<String, FieldStatsInfo>() {
            {
                put("score", level1Pivot1FieldStatsInfo);
            }
        });
        when(level1Pivot2.getValue()).thenReturn("field1value2");
        when(level1Pivot2.getCount()).thenReturn(2);
        when(level1Pivot2FieldStatsInfo.getSum()).thenReturn(2.0);
        when(level1Pivot2.getFieldStatsInfo()).thenReturn(new HashMap<String, FieldStatsInfo>() {
            {
                put("score", level1Pivot2FieldStatsInfo);
            }
        });
        when(level2Pivot1.getValue()).thenReturn("field2value1");
        when(level2Pivot1.getCount()).thenReturn(3);
        when(level2Pivot1FieldStatsInfo.getSum()).thenReturn(3.0);
        when(level2Pivot1.getFieldStatsInfo()).thenReturn(new HashMap<String, FieldStatsInfo>() {
            {
                put("score", level2Pivot1FieldStatsInfo);
            }
        });
        when(level2Pivot2.getValue()).thenReturn("field2value2");
        when(level2Pivot2.getCount()).thenReturn(4);
        when(level2Pivot2FieldStatsInfo.getSum()).thenReturn(4.0);
        when(level2Pivot2.getFieldStatsInfo()).thenReturn(new HashMap<String, FieldStatsInfo>() {
            {
                put("score", level2Pivot2FieldStatsInfo);
            }
        });
        when(level1Pivot1.getPivot()).thenReturn(level2Pivots);

        List<GroupResult> level1GroupResults = solrSearchDao.getGroupResults(groupRequest, 0, level1Pivots);

        assertEquals("field1value1", level1GroupResults.get(0).getKey());
        assertEquals(1, level1GroupResults.get(0).getTotal());
        assertEquals(1.0, level1GroupResults.get(0).getScore(), 0.00001);
        assertEquals("field2", level1GroupResults.get(0).getGroupedBy());
        assertEquals("field1value2", level1GroupResults.get(1).getKey());
        assertEquals(2, level1GroupResults.get(1).getTotal());
        assertEquals(2.0, level1GroupResults.get(1).getScore(), 0.00001);
        assertEquals("field2", level1GroupResults.get(1).getGroupedBy());
        assertEquals(0, level1GroupResults.get(1).getGroupResults().size());

        List<GroupResult> level2GroupResults = level1GroupResults.get(0).getGroupResults();
        assertEquals("field2value2", level2GroupResults.get(0).getKey());
        assertEquals(4, level2GroupResults.get(0).getTotal());
        assertEquals(4.0, level2GroupResults.get(0).getScore(), 0.00001);
        assertNull(level2GroupResults.get(0).getGroupedBy());
        assertNull(level2GroupResults.get(0).getGroupResults());
        assertEquals("field2value1", level2GroupResults.get(1).getKey());
        assertEquals(3, level2GroupResults.get(1).getTotal());
        assertEquals(3.0, level2GroupResults.get(1).getScore(), 0.00001);
        assertNull(level2GroupResults.get(1).getGroupedBy());
        assertNull(level2GroupResults.get(1).getGroupResults());
    }

}