Separate QuantizationStateCache teardown from KNNTestCase

Signed-off-by: owenhalpert <[email protected]>
owenhalpert · Jan 18, 2025 · 8b14477 · 8b14477
1 parent eecb45c
commit 8b14477
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 7 deletions.
diff --git a/src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java b/src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java
@@ -35,7 +35,6 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.index.IndexService;
-import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.test.OpenSearchSingleNodeTestCase;
@@ -87,7 +86,6 @@ protected boolean resetNodeAfterTest() {
     public void tearDown() throws Exception {
         NativeMemoryCacheManager.getInstance().invalidateAll();
         NativeMemoryCacheManager.getInstance().close();
-        QuantizationStateCacheManager.getInstance().close();
         NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance().close();
         NativeMemoryLoadStrategy.TrainingLoadStrategy.getInstance().close();
         NativeMemoryLoadStrategy.AnonymousLoadStrategy.getInstance().close();

diff --git a/src/test/java/org/opensearch/knn/KNNTestCase.java b/src/test/java/org/opensearch/knn/KNNTestCase.java
@@ -24,10 +24,8 @@
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.common.xcontent.XContentHelper;
-import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
 import org.opensearch.test.OpenSearchTestCase;
 
-import java.io.IOException;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Map;
@@ -75,7 +73,7 @@ protected boolean enableWarningsCheck() {
         return false;
     }
 
-    public void resetState() throws IOException {
+    public void resetState() {
         // Reset all of the counters
         for (KNNCounter knnCounter : KNNCounter.values()) {
             knnCounter.set(0L);
@@ -85,7 +83,6 @@ public void resetState() throws IOException {
         // Clean up the cache
         NativeMemoryCacheManager.getInstance().invalidateAll();
         NativeMemoryCacheManager.getInstance().close();
-        QuantizationStateCacheManager.getInstance().close();
     }
 
     private void initKNNSettings() {

diff --git a/...org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCacheTests.java b/...org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCacheTests.java
@@ -42,8 +42,9 @@ public void setThreadPool() {
     }
 
     @After
-    public void terminateThreadPool() {
+    public void stopMaintenance() throws IOException {
         terminate(threadPool);
+        QuantizationStateCache.getInstance().close();
     }
 
     @SneakyThrows