Fix tests

Signed-off-by: Rishab Nahata <[email protected]>

Author: imRishN

server/src/test/java/org/opensearch/search/aggregations/startree/BoolStarTreeFilterProviderTests.java

@@ -11,6 +11,7 @@
 import org.apache.lucene.document.FloatPoint;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.NumericUtils;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.compositeindex.datacube.Metric;
 import org.opensearch.index.compositeindex.datacube.MetricStat;
 import org.opensearch.index.compositeindex.datacube.OrdinalDimension;
@@ -34,6 +35,7 @@
 import org.opensearch.search.startree.filter.StarTreeFilter;
 import org.opensearch.search.startree.filter.provider.StarTreeFilterProvider;
 import org.opensearch.test.OpenSearchTestCase;
+import org.junit.After;
 import org.junit.Before;
 
 import java.io.IOException;
@@ -44,13 +46,15 @@
 import java.util.Optional;
 import java.util.Set;
 
+import static org.opensearch.common.util.FeatureFlags.STAR_TREE_INDEX;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 public class BoolStarTreeFilterProviderTests extends OpenSearchTestCase {
     private SearchContext searchContext;
     private MapperService mapperService;
     private CompositeDataCubeFieldType compositeFieldType;
+    private static FeatureFlags.TestUtils.FlagWriteLock fflock = null;
 
     @Before
     public void setup() {
@@ -106,6 +110,12 @@ public void setup() {
                 )
             )
         );
+        fflock = new FeatureFlags.TestUtils.FlagWriteLock(STAR_TREE_INDEX);
+    }
+
+    @After
+    public void teardown() throws IOException {
+        fflock.close();
     }
 
     public void testSimpleMustWithMultipleDimensions() throws IOException {

server/src/test/java/org/opensearch/search/aggregations/startree/DimensionFilterAndMapperTests.java

@@ -10,6 +10,7 @@
 
 import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.util.BytesRef;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.compositeindex.datacube.Metric;
 import org.opensearch.index.compositeindex.datacube.MetricStat;
 import org.opensearch.index.compositeindex.datacube.OrdinalDimension;
@@ -36,17 +37,32 @@
 import org.opensearch.search.startree.filter.provider.DimensionFilterMapper;
 import org.opensearch.search.startree.filter.provider.StarTreeFilterProvider;
 import org.opensearch.test.OpenSearchTestCase;
+import org.junit.After;
+import org.junit.Before;
 
 import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
+import static org.opensearch.common.util.FeatureFlags.STAR_TREE_INDEX;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 public class DimensionFilterAndMapperTests extends OpenSearchTestCase {
 
+    private static FeatureFlags.TestUtils.FlagWriteLock fflock = null;
+
+    @Before
+    public void setup() {
+        fflock = new FeatureFlags.TestUtils.FlagWriteLock(STAR_TREE_INDEX);
+    }
+
+    @After
+    public void teardown() throws IOException {
+        fflock.close();
+    }
+
     public void testKeywordOrdinalMapping() throws IOException {
         DimensionFilterMapper dimensionFilterMapper = DimensionFilterMapper.Factory.fromMappedFieldType(
             new KeywordFieldMapper.KeywordFieldType("keyword")

server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java

@@ -34,6 +34,7 @@
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.opensearch.common.lucene.Lucene;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.util.MockBigArrays;
 import org.opensearch.common.util.MockPageCacheRecycler;
 import org.opensearch.core.indices.breaker.CircuitBreakerService;
@@ -74,6 +75,8 @@
 import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
 import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
 import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
+import org.junit.After;
+import org.junit.Before;
 
 import java.io.IOException;
 import java.math.BigInteger;
@@ -90,6 +93,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static org.opensearch.common.util.FeatureFlags.STAR_TREE_INDEX;
 import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.objectToUnsignedLong;
 import static org.opensearch.search.aggregations.AggregationBuilders.avg;
 import static org.opensearch.search.aggregations.AggregationBuilders.count;
@@ -104,6 +108,17 @@ public class MetricAggregatorTests extends AggregatorTestCase {
     private static final String FIELD_NAME = "field";
     private static final NumberFieldMapper.NumberType DEFAULT_FIELD_TYPE = NumberFieldMapper.NumberType.LONG;
     private static final MappedFieldType DEFAULT_MAPPED_FIELD = new NumberFieldMapper.NumberFieldType(FIELD_NAME, DEFAULT_FIELD_TYPE);
+    private static FeatureFlags.TestUtils.FlagWriteLock fflock = null;
+
+    @Before
+    public void setup() {
+        fflock = new FeatureFlags.TestUtils.FlagWriteLock(STAR_TREE_INDEX);
+    }
+
+    @After
+    public void teardown() throws IOException {
+        fflock.close();
+    }
 
     protected Codec getCodec(
         Supplier<Integer> maxLeafDocsSupplier,