FIX: Added resetState for uTs so state does not spill over (#159)

opendistro-for-elasticsearch · Aug 12, 2020 · 73f2896 · 73f2896
1 parent f345ac1
commit 73f2896
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 15 deletions.
diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/knn/KNNTestCase.java b/src/test/java/com/amazon/opendistroforelasticsearch/knn/KNNTestCase.java
@@ -15,19 +15,28 @@
 
 package com.amazon.opendistroforelasticsearch.knn;
 
+import com.amazon.opendistroforelasticsearch.knn.index.KNNIndexCache;
 import com.amazon.opendistroforelasticsearch.knn.plugin.stats.KNNCounter;
 import org.elasticsearch.test.ESTestCase;
-import org.junit.Before;
 
 /**
  * Base class for integration tests for KNN plugin. Contains several methods for testing KNN ES functionality.
  */
 public class KNNTestCase extends ESTestCase {
-    @Before
-    public void resetCounters() {
+    @Override
+    public void tearDown() throws Exception {
+        super.tearDown();
+        resetState();
+    }
+
+    public static void resetState() {
         // Reset all of the counters
         for (KNNCounter knnCounter : KNNCounter.values()) {
             knnCounter.set(0L);
         }
+
+        // Clean up the cache
+        KNNIndexCache.getInstance().evictAllGraphsFromCache();
+        KNNIndexCache.getInstance().close();
     }
 }
diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/knn/index/KNNIndexCacheTests.java b/src/test/java/com/amazon/opendistroforelasticsearch/knn/index/KNNIndexCacheTests.java
@@ -15,8 +15,8 @@
 
 package com.amazon.opendistroforelasticsearch.knn.index;
 
+import com.amazon.opendistroforelasticsearch.knn.KNNTestCase;
 import com.amazon.opendistroforelasticsearch.knn.plugin.KNNPlugin;
-import com.amazon.opendistroforelasticsearch.knn.plugin.stats.KNNCounter;
 import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.index.IndexResponse;
@@ -42,15 +42,6 @@ public class KNNIndexCacheTests extends ESSingleNodeTestCase {
     private final String testIndexName = "test_index";
     private final String testFieldName = "test_field";
 
-    @Override
-    public void setUp() throws Exception {
-        // Reset all of the counters
-        for (KNNCounter knnCounter : KNNCounter.values()) {
-            knnCounter.set(0L);
-        }
-        super.setUp();
-    }
-
     @Override
     protected Collection<Class<? extends Plugin>> getPlugins() {
         return Collections.singletonList(KNNPlugin.class);
@@ -64,8 +55,7 @@ protected boolean resetNodeAfterTest() {
     @Override
     public void tearDown() throws Exception {
         super.tearDown();
-        KNNIndexCache.getInstance().evictAllGraphsFromCache();
-        KNNIndexCache.getInstance().close();
+        KNNTestCase.resetState();
     }
 
     public void testGetIndicesCacheStats() throws IOException, InterruptedException, ExecutionException {