improve the unit test coverage for query insight plugin

Signed-off-by: Chenyang Ji <cyji@amazon.com>

plugins/query-insights/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryLatencyRecord.java

@@ -69,4 +69,19 @@ public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeMap(phaseLatencyMap, StreamOutput::writeString, StreamOutput::writeLong);
     }
+
+    public boolean equals(SearchQueryLatencyRecord other) {
+        if (!super.equals(other)) {
+            return false;
+        }
+        for (String key : phaseLatencyMap.keySet()) {
+            if (!other.getPhaseLatencyMap().containsKey(key)) {
+                return false;
+            }
+            if(!phaseLatencyMap.get(key).equals(other.getPhaseLatencyMap().get(key))) {
+                return false;
+            }
+        }
+        return true;
+    }
 }

plugins/query-insights/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java

@@ -158,6 +158,32 @@ public int compareTo(SearchQueryRecord<T> otherRecord) {
         return value.compareTo(otherRecord.getValue());
     }
 
+    public boolean equals(SearchQueryRecord<T> other) {
+        if(false == this.timestamp.equals(other.getTimestamp()) &&
+            this.searchType.equals(other.getSearchType()) &&
+            this.source.equals(other.getSource()) &&
+            this.totalShards == other.getTotalShards() &&
+            this.indices.length == other.getIndices().length &&
+            this.propertyMap.size() == other.getPropertyMap().size() &&
+            this.value.equals(other.getValue())) {
+            return false;
+        }
+        for (int i = 0; i < indices.length; i++) {
+            if (!indices[i].equals(other.getIndices()[i])) {
+                return false;
+            }
+        }
+        for (String key : propertyMap.keySet()) {
+            if (!other.getPropertyMap().containsKey(key)) {
+                return false;
+            }
+            if(!propertyMap.get(key).equals(other.getPropertyMap().get(key))) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     @SuppressWarnings("unchecked")
     private T castToValue(Object obj) throws ClassCastException {
         try {

plugins/query-insights/src/test/java/org/opensearch/plugin/insights/QueryInsightTestUtils.java

@@ -0,0 +1,167 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.insights;
+
+import org.opensearch.action.search.SearchType;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.plugin.insights.rules.action.top_queries.TopQueries;
+import org.opensearch.plugin.insights.rules.model.SearchQueryLatencyRecord;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.test.VersionUtils;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeSet;
+
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.emptySet;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
+import static org.opensearch.test.OpenSearchTestCase.buildNewFakeTransportAddress;
+import static org.opensearch.test.OpenSearchTestCase.random;
+import static org.opensearch.test.OpenSearchTestCase.randomAlphaOfLengthBetween;
+import static org.opensearch.test.OpenSearchTestCase.randomArray;
+import static org.opensearch.test.OpenSearchTestCase.randomInt;
+import static org.opensearch.test.OpenSearchTestCase.randomIntBetween;
+import static org.opensearch.test.OpenSearchTestCase.randomLong;
+
+final public class QueryInsightTestUtils {
+
+    public QueryInsightTestUtils() {}
+
+    public static List<SearchQueryLatencyRecord> generateQueryInsightRecords(int count) {
+        return generateQueryInsightRecords(count, count);
+    }
+
+    /**
+     * Creates a List of random Query Insight Records for testing purpose
+     */
+    public static List<SearchQueryLatencyRecord> generateQueryInsightRecords(int lower, int upper) {
+        List<SearchQueryLatencyRecord> records = new ArrayList<>();
+        int countOfRecords = randomIntBetween(lower, upper);
+        for (int i = 0; i < countOfRecords; ++i) {
+            Map<String, Object> propertyMap = new HashMap<>();
+            int countOfProperties = randomIntBetween(2, 5);
+            for (int j = 0; j < countOfProperties; ++j) {
+                propertyMap.put(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLengthBetween(5, 10));
+            }
+            Map<String, Long> phaseLatencyMap = new HashMap<>();
+            int countOfPhases = randomIntBetween(2, 5);
+            for (int j = 0; j < countOfPhases; ++j) {
+                phaseLatencyMap.put(randomAlphaOfLengthBetween(5, 10), randomLong());
+            }
+            records.add(new SearchQueryLatencyRecord(
+                System.currentTimeMillis(),
+                SearchType.QUERY_THEN_FETCH,
+                "{\"size\":20}",
+                randomIntBetween(1, 100),
+                randomArray(1, 3, String[]::new, () -> randomAlphaOfLengthBetween(5, 10)),
+                propertyMap,
+                phaseLatencyMap
+            ));
+        }
+        return records;
+    }
+
+    public static TopQueries createTopQueries() {
+        DiscoveryNode node = new DiscoveryNode(
+            "node_for_top_queries_test",
+            buildNewFakeTransportAddress(),
+            emptyMap(),
+            emptySet(),
+            VersionUtils.randomVersion(random())
+        );
+
+        Map<String, Object> propertyMap = new HashMap<>();
+        propertyMap.put("userId", "user1");
+
+        Map<String, Long> phaseLatencyMap = new HashMap<>();
+        phaseLatencyMap.put("expand", 0L);
+        phaseLatencyMap.put("query", 20L);
+        phaseLatencyMap.put("fetch", 1L);
+
+        List<SearchQueryLatencyRecord> records = new ArrayList<>();
+        records.add(new SearchQueryLatencyRecord(
+            randomLong(),
+            SearchType.QUERY_THEN_FETCH,
+            "{\"size\":20}",
+            randomInt(),
+            randomArray(1, 3, String[]::new, () -> randomAlphaOfLengthBetween(5, 10)),
+            propertyMap,
+            phaseLatencyMap
+        ));
+
+        return new TopQueries(node, records);
+    }
+
+    public static void compareJson(ToXContent param1, ToXContent param2) throws IOException {
+        if (param1 == null || param2 == null) {
+            assertNull(param1);
+            assertNull(param2);
+            return;
+        }
+
+        ToXContent.Params params = ToXContent.EMPTY_PARAMS;
+        XContentBuilder param1Builder = jsonBuilder();
+        param1.toXContent(param1Builder, params);
+
+        XContentBuilder param2Builder = jsonBuilder();
+        param2.toXContent(param2Builder, params);
+
+        assertEquals(param1Builder.toString(), param2Builder.toString());
+    }
+
+    public static boolean checkRecordsEquals(List<SearchQueryLatencyRecord> records1, List<SearchQueryLatencyRecord> records2) {
+        if (records1.size() != records2.size()) {
+            return false;
+        }
+        for (int i = 0; i < records1.size(); i++) {
+            if (!records1.get(i).equals(records2.get(i))) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+
+
+    public static boolean checkRecordsEqualsWithoutOrder(List<SearchQueryLatencyRecord> records1, List<SearchQueryLatencyRecord> records2) {
+        Set<SearchQueryLatencyRecord> set2 = new TreeSet<>(new LatencyRecordComparator());
+        set2.addAll(records2);
+        if (records1.size() != records2.size()) {
+            return false;
+        }
+        for (int i = 0; i < records1.size(); i++) {
+            if (!set2.contains(records1.get(i))) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    static class LatencyRecordComparator implements Comparator<SearchQueryLatencyRecord> {
+        @Override
+        public int compare(SearchQueryLatencyRecord record1, SearchQueryLatencyRecord record2)
+        {
+            if (record1.equals(record2)) {
+                return 0;
+            }
+            return record1.compareTo(record2);
+        }
+    }
+}