Refactor unit tests

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
opensearch-project · Aug 5, 2024 · f83448c · f83448c
1 parent 1fd5717
commit f83448c
Show file tree

Hide file tree

Showing 2 changed files with 122 additions and 444 deletions.
diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java
@@ -178,6 +178,24 @@ public static List<SearchQueryRecord> generateQueryInsightsRecordsWithMeasuremen
         return records;
     }
 
+    public static List<List<SearchQueryRecord>> generateMultipleQueryInsightsRecordsWithMeasurement(
+        int count,
+        MetricType metricType,
+        List<Number> measurements
+    ) {
+        List<List<SearchQueryRecord>> multipleRecordLists = new ArrayList<>();
+
+        for (int i = 0; i < measurements.size(); i++) {
+            List<SearchQueryRecord> records = generateQueryInsightRecords(count);
+            multipleRecordLists.add(records);
+            for (SearchQueryRecord record : records) {
+                record.getMeasurements().get(metricType).setMeasurement(measurements.get(i));
+            }
+            QueryInsightsTestUtils.populateHashcode(records, i);
+        }
+        return multipleRecordLists;
+    }
+
     public static void populateSameQueryHashcodes(List<SearchQueryRecord> searchQueryRecords) {
         for (SearchQueryRecord record : searchQueryRecords) {
             record.getAttributes().put(Attribute.QUERY_HASHCODE, 1);