update warmup rolling upgrade scenario

Signed-off-by: Vijayan Balasubramanian <[email protected]>
opensearch-project · Oct 24, 2024 · 416816b · 416816b
1 parent 5bbd300
commit 416816b
Showing 1 changed file with 14 additions and 32 deletions.
diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java
@@ -5,7 +5,11 @@
 
 package org.opensearch.knn.bwc;
 
+import org.opensearch.common.settings.Settings;
+import org.opensearch.knn.index.KNNSettings;
+
 import java.util.Collections;
+import java.util.List;
 
 import static org.opensearch.knn.TestUtils.NODES_BWC_CLUSTER;
 
@@ -20,44 +24,22 @@ public void testKNNWarmup() throws Exception {
         switch (getClusterType()) {
             case OLD:
                 createKnnIndex(testIndex, getKNNDefaultIndexSettings(), createKnnIndexMapping(TEST_FIELD, DIMENSIONS));
-                int docIdOld = 0;
-                addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, docIdOld, NUM_DOCS);
+                addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, 0, NUM_DOCS);
                 break;
             case MIXED:
-                int totalDocsCountMixed;
-                int docIdMixed;
-                if (isFirstMixedRound()) {
-                    docIdMixed = NUM_DOCS;
-                    totalDocsCountMixed = 2 * NUM_DOCS;
-                } else {
-                    docIdMixed = 2 * NUM_DOCS;
-                    totalDocsCountMixed = 3 * NUM_DOCS;
-                }
-                validateKNNWarmupOnUpgrade(totalDocsCountMixed, docIdMixed);
+                int graphCount = getTotalGraphsInCache();
+                knnWarmup(Collections.singletonList(testIndex));
+                assertTrue(getTotalGraphsInCache() > graphCount);
+                clearCache(List.of(testIndex));
                 break;
             case UPGRADED:
-                int docIdUpgraded = 3 * NUM_DOCS;
-                int totalDocsCountUpgraded = 4 * NUM_DOCS;
-                validateKNNWarmupOnUpgrade(totalDocsCountUpgraded, docIdUpgraded);
-
+                updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0));
+                addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, NUM_DOCS, NUM_DOCS);
+                int updatedGraphCount = getTotalGraphsInCache();
+                knnWarmup(Collections.singletonList(testIndex));
+                assertTrue(getTotalGraphsInCache() > updatedGraphCount);
                 deleteKNNIndex(testIndex);
         }
-
-    }
-
-    // validation steps for KNN Warmup after upgrading each node from old version to new version
-    public void validateKNNWarmupOnUpgrade(int totalDocsCount, int docId) throws Exception {
-        int graphCount = getTotalGraphsInCache();
-        knnWarmup(Collections.singletonList(testIndex));
-        assertTrue(getTotalGraphsInCache() > graphCount);
-
-        addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, docId, NUM_DOCS);
-
-        int updatedGraphCount = getTotalGraphsInCache();
-        knnWarmup(Collections.singletonList(testIndex));
-        assertTrue(getTotalGraphsInCache() > updatedGraphCount);
-
-        validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, totalDocsCount, K);
     }
 
 }