Make random sampling method more dynamic

Author: gmarouli

modules/data-streams/src/test/java/org/elasticsearch/datastreams/lifecycle/DataStreamLifecycleFixtures.java

@@ -40,6 +40,7 @@
 import static org.elasticsearch.test.ESIntegTestCase.client;
 import static org.elasticsearch.test.ESTestCase.between;
 import static org.elasticsearch.test.ESTestCase.frequently;
+import static org.elasticsearch.test.ESTestCase.randomFrom;
 import static org.elasticsearch.test.ESTestCase.randomIntBetween;
 import static org.junit.Assert.assertTrue;
 
@@ -184,11 +185,10 @@ private static DataStreamLifecycle.DownsamplingRound nextRound(DataStreamLifecyc
     }
 
     public static DownsampleConfig.SamplingMethod randomSamplingMethod() {
-        return switch (between(0, 2)) {
-            case 0 -> null;
-            case 1 -> DownsampleConfig.SamplingMethod.AGGREGATE;
-            case 2 -> DownsampleConfig.SamplingMethod.LAST_VALUE;
-            default -> throw new IllegalStateException("Unknown randomisation path");
-        };
+        if (between(0, DownsampleConfig.SamplingMethod.values().length) == 0) {
+            return null;
+        } else {
+            return randomFrom(DownsampleConfig.SamplingMethod.values());
+        }
     }
 }

server/src/test/java/org/elasticsearch/action/downsample/DownsampleConfigTests.java

@@ -50,12 +50,11 @@ public static DownsampleConfig randomConfig() {
     }
 
     public static DownsampleConfig.SamplingMethod randomSamplingMethod() {
-        return switch (between(0, 2)) {
-            case 0 -> null;
-            case 1 -> DownsampleConfig.SamplingMethod.AGGREGATE;
-            case 2 -> DownsampleConfig.SamplingMethod.LAST_VALUE;
-            default -> throw new AssertionError("Illegal randomisation branch");
-        };
+        if (between(0, DownsampleConfig.SamplingMethod.values().length) == 0) {
+            return null;
+        } else {
+            return randomFrom(DownsampleConfig.SamplingMethod.values());
+        }
     }
 
     @Override

x-pack/plugin/downsample/src/internalClusterTest/java/org/elasticsearch/xpack/downsample/DownsamplingIntegTestCase.java

@@ -336,11 +336,10 @@ private static boolean hasTimeSeriesDimensionTrue(Map<String, ?> fieldMapping) {
     }
 
     public static DownsampleConfig.SamplingMethod randomSamplingMethod() {
-        return switch (between(0, 2)) {
-            case 0 -> null;
-            case 1 -> DownsampleConfig.SamplingMethod.AGGREGATE;
-            case 2 -> DownsampleConfig.SamplingMethod.LAST_VALUE;
-            default -> throw new IllegalStateException("Unexpected randomisation branch");
-        };
+        if (between(0, DownsampleConfig.SamplingMethod.values().length) == 0) {
+            return null;
+        } else {
+            return randomFrom(DownsampleConfig.SamplingMethod.values());
+        }
     }
 }

x-pack/plugin/downsample/src/test/java/org/elasticsearch/xpack/downsample/DownsampleActionSingleNodeTests.java

@@ -1824,11 +1824,10 @@ public void testDuplicateDownsampleRequest() throws Exception {
     }
 
     static DownsampleConfig.SamplingMethod randomSamplingMethod() {
-        return switch (between(0, 2)) {
-            case 0 -> null;
-            case 1 -> DownsampleConfig.SamplingMethod.AGGREGATE;
-            case 2 -> DownsampleConfig.SamplingMethod.LAST_VALUE;
-            default -> throw new IllegalStateException("Unexpected randomisation branch");
-        };
+        if (between(0, DownsampleConfig.SamplingMethod.values().length) == 0) {
+            return null;
+        } else {
+            return randomFrom(DownsampleConfig.SamplingMethod.values());
+        }
     }
 }

x-pack/plugin/security/src/internalClusterTest/java/org/elasticsearch/integration/DataStreamLifecycleDownsamplingSecurityIT.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.action.datastreams.GetDataStreamAction;
 import org.elasticsearch.action.datastreams.lifecycle.ErrorEntry;
+import org.elasticsearch.action.downsample.DownsampleConfig;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.client.internal.Client;
@@ -120,6 +121,7 @@ public void testDownsamplingAuthorized() throws Exception {
         String dataStreamName = "metrics-foo";
 
         DataStreamLifecycle.Template lifecycle = DataStreamLifecycle.dataLifecycleBuilder()
+            .downsamplingMethod(randomSamplingMethod())
             .downsamplingRounds(
                 List.of(
                     new DataStreamLifecycle.DownsamplingRound(TimeValue.timeValueMillis(0), new DateHistogramInterval("5m")),
@@ -401,9 +403,18 @@ private void bulkIndex(Client client, String dataStreamName, Supplier<XContentBu
         logger.info("-> Indexed [{}] documents. Dropped [{}] duplicates.", docsIndexed, duplicates);
     }
 
+    private static DownsampleConfig.SamplingMethod randomSamplingMethod() {
+        if (between(0, DownsampleConfig.SamplingMethod.values().length) == 0) {
+            return null;
+        } else {
+            return randomFrom(DownsampleConfig.SamplingMethod.values());
+        }
+    }
+
     public static class SystemDataStreamWithDownsamplingConfigurationPlugin extends Plugin implements SystemIndexPlugin {
 
         public static final DataStreamLifecycle.Template LIFECYCLE = DataStreamLifecycle.dataLifecycleBuilder()
+            .downsamplingMethod(randomSamplingMethod())
             .downsamplingRounds(
                 List.of(
                     new DataStreamLifecycle.DownsamplingRound(TimeValue.timeValueMillis(0), new DateHistogramInterval("5m")),