Add block loader tests for aggregate_metric_double

Author: lkts

test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java

@@ -16,6 +16,7 @@
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 
 import static org.junit.Assert.assertNotNull;
@@ -292,12 +293,29 @@ public BlockLoader.IntBuilder count() {
 
         @Override
         public BlockLoader.Block build() {
-            var minBlock = max.build();
+            var minBlock = min.build();
             var maxBlock = max.build();
-            var sumBlock = max.build();
-            var countBlock = max.build();
+            var sumBlock = sum.build();
+            var countBlock = count.build();
 
-            return new TestBlock(null);
+            assert minBlock.size() == maxBlock.size();
+            assert maxBlock.size() == sumBlock.size();
+            assert sumBlock.size() == countBlock.size();
+
+            var values = new ArrayList<>(minBlock.size());
+
+            for (int i = 0; i < minBlock.size(); i++) {
+                // we need to represent this complex block somehow
+                var value = new HashMap<String, Object>();
+                value.put("min", minBlock.values.get(i));
+                value.put("max", maxBlock.values.get(i));
+                value.put("sum", sumBlock.values.get(i));
+                value.put("value_count", countBlock.values.get(i));
+
+                values.add(value);
+            }
+
+            return new TestBlock(values);
         }
 
         @Override

test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceHandler.java

@@ -82,6 +82,10 @@ default DataSourceResponse.VersionStringGenerator handle(DataSourceRequest.Versi
         return null;
     }
 
+    default DataSourceResponse.AggregateMetricDoubleGenerator handle(DataSourceRequest.AggregateMetricDoubleGenerator request) {
+        return null;
+    }
+
     default DataSourceResponse.NullWrapper handle(DataSourceRequest.NullWrapper request) {
         return null;
     }

test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceRequest.java

@@ -132,6 +132,12 @@ public DataSourceResponse.VersionStringGenerator accept(DataSourceHandler handle
         }
     }
 
+    record AggregateMetricDoubleGenerator() implements DataSourceRequest<DataSourceResponse.AggregateMetricDoubleGenerator> {
+        public DataSourceResponse.AggregateMetricDoubleGenerator accept(DataSourceHandler handler) {
+            return handler.handle(this);
+        }
+    }
+
     record NullWrapper() implements DataSourceRequest<DataSourceResponse.NullWrapper> {
         public DataSourceResponse.NullWrapper accept(DataSourceHandler handler) {
             return handler.handle(this);

test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/datasource/DataSourceResponse.java

@@ -55,6 +55,8 @@ record IpGenerator(Supplier<InetAddress> generator) implements DataSourceRespons
 
     record VersionStringGenerator(Supplier<String> generator) implements DataSourceResponse {}
 
+    record AggregateMetricDoubleGenerator(Supplier<Map<String, Number>> generator) implements DataSourceResponse {}
+
     record NullWrapper(Function<Supplier<Object>, Supplier<Object>> wrapper) implements DataSourceResponse {}
 
     record ArrayWrapper(Function<Supplier<Object>, Supplier<Object>> wrapper) implements DataSourceResponse {}

x-pack/plugin/mapper-aggregate-metric/src/main/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleFieldMapper.java

@@ -851,6 +851,7 @@ protected void parseCreateField(DocumentParserContext context) throws IOExceptio
             // by its FieldMapper#parse()
             throw e;
         }
+
         for (Map.Entry<Metric, Number> parsed : metricsParsed.entrySet()) {
             NumberFieldMapper delegateFieldMapper = metricFieldMappers.get(parsed.getKey());
             delegateFieldMapper.indexValue(context, parsed.getValue());

x-pack/plugin/mapper-aggregate-metric/src/test/java/org/elasticsearch/xpack/aggregatemetric/mapper/AggregateMetricDoubleFieldBlockLoaderTests.java

@@ -8,21 +8,56 @@
 package org.elasticsearch.xpack.aggregatemetric.mapper;
 
 import org.elasticsearch.index.mapper.BlockLoaderTestCase;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceHandler;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceRequest;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceResponse;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.xpack.aggregatemetric.AggregateMetricMapperPlugin;
+import org.elasticsearch.xpack.aggregatemetric.mapper.datageneration.AggregateMetricDoubleDataSourceHandler;
 
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 public class AggregateMetricDoubleFieldBlockLoaderTests extends BlockLoaderTestCase {
     public AggregateMetricDoubleFieldBlockLoaderTests(Params params) {
-        super("aggregate_metric_double", params);
+        super("aggregate_metric_double", List.of(new AggregateMetricDoubleDataSourceHandler(), new DataSourceHandler() {
+            @Override
+            public DataSourceResponse.ObjectArrayGenerator handle(DataSourceRequest.ObjectArrayGenerator request) {
+                // aggregate_metric_double does not support multiple values in a document so we can't have object arrays
+                return new DataSourceResponse.ObjectArrayGenerator(Optional::empty);
+            }
+        }), params);
     }
 
     @Override
     protected Object expected(Map<String, Object> fieldMapping, Object value, TestContext testContext) {
-        return null;
+        if (value instanceof Map<?, ?> map) {
+            var expected = new HashMap<String, Object>(map.size());
+
+            // put explicit `null` for metrics that are not present, this is how the block looks like
+            Arrays.stream(AggregateMetricDoubleFieldMapper.Metric.values())
+                .map(AggregateMetricDoubleFieldMapper.Metric::toString)
+                .sorted()
+                .forEach(m -> {
+                    expected.put(m, map.get(m));
+                });
+
+            return expected;
+        }
+
+        // malformed or null, return "empty row"
+        return new HashMap<>() {
+            {
+                put("min", null);
+                put("max", null);
+                put("sum", null);
+                put("value_count", null);
+            }
+        };
     }
 
     @Override

x-pack/plugin/mapper-aggregate-metric/src/test/java/org/elasticsearch/xpack/aggregatemetric/mapper/datageneration/AggregateMetricDoubleDataSourceHandler.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.aggregatemetric.mapper.datageneration;
+
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceHandler;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceRequest;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceResponse;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.aggregatemetric.mapper.AggregateMetricDoubleFieldMapper;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+
+public class AggregateMetricDoubleDataSourceHandler implements DataSourceHandler {
+    @Override
+    public DataSourceResponse.AggregateMetricDoubleGenerator handle(DataSourceRequest.AggregateMetricDoubleGenerator request) {
+        return new DataSourceResponse.AggregateMetricDoubleGenerator(() -> {
+            var metricContainer = new HashMap<String, Number>();
+
+            // min and max must make sense - max has to be gte min
+            double min = ESTestCase.randomDoubleBetween(-Double.MAX_VALUE, 1_000_000_000, false);
+            double max = ESTestCase.randomDoubleBetween(min, Double.MAX_VALUE, true);
+
+            metricContainer.put("min", min);
+            metricContainer.put("max", max);
+            metricContainer.put("sum", ESTestCase.randomDouble());
+            metricContainer.put("value_count", ESTestCase.randomIntBetween(1, Integer.MAX_VALUE));
+
+            return metricContainer;
+        });
+    }
+
+    @Override
+    public DataSourceResponse.LeafMappingParametersGenerator handle(DataSourceRequest.LeafMappingParametersGenerator request) {
+        if (request.fieldType().equals("aggregate_metric_double") == false) {
+            return null;
+        }
+
+        return new DataSourceResponse.LeafMappingParametersGenerator(() -> {
+            var map = new HashMap<String, Object>();
+
+            List<AggregateMetricDoubleFieldMapper.Metric> metrics = ESTestCase.randomNonEmptySubsetOf(
+                Arrays.asList(AggregateMetricDoubleFieldMapper.Metric.values())
+            );
+
+            map.put("metrics", metrics.stream().map(Enum::toString).toList());
+            map.put("default_metric", metrics.get(ESTestCase.randomIntBetween(0, metrics.size() - 1)));
+
+            if (ESTestCase.randomBoolean()) {
+                map.put("ignore_malformed", ESTestCase.randomBoolean());
+            }
+
+            return map;
+        });
+    }
+
+    @Override
+    public DataSourceResponse.FieldDataGenerator handle(DataSourceRequest.FieldDataGenerator request) {
+        if (request.fieldType().equals("aggregate_metric_double") == false) {
+            return null;
+        }
+
+        return new DataSourceResponse.FieldDataGenerator(new AggregateMetricDoubleFieldDataGenerator(request.dataSource()));
+    }
+}

x-pack/plugin/mapper-aggregate-metric/src/test/java/org/elasticsearch/xpack/aggregatemetric/mapper/datageneration/AggregateMetricDoubleFieldDataGenerator.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.aggregatemetric.mapper.datageneration;
+
+import org.elasticsearch.logsdb.datageneration.FieldDataGenerator;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSource;
+import org.elasticsearch.logsdb.datageneration.datasource.DataSourceRequest;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+public class AggregateMetricDoubleFieldDataGenerator implements FieldDataGenerator {
+    private final Supplier<Map<String, Number>> metrics;
+    private final Function<Supplier<Object>, Supplier<Object>> nullWrapper;
+    private final Function<Supplier<Object>, Supplier<Object>> nullAndMalformedWrapper;
+
+    public AggregateMetricDoubleFieldDataGenerator(DataSource dataSource) {
+        this.metrics = dataSource.get(new DataSourceRequest.AggregateMetricDoubleGenerator()).generator();
+
+        this.nullWrapper = dataSource.get(new DataSourceRequest.NullWrapper()).wrapper();
+
+        var strings = dataSource.get(new DataSourceRequest.StringGenerator()).generator();
+        var malformed = dataSource.get(new DataSourceRequest.MalformedWrapper(strings::get)).wrapper();
+        this.nullAndMalformedWrapper = malformed.andThen(nullWrapper);
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public Object generateValue(Map<String, Object> fieldMapping) {
+        if (fieldMapping == null) {
+            // this field can't be properly mapped with dynamic mapping
+            return null;
+        }
+
+        // metrics returned have all metric fields but we only need those that appear in the mapping
+        Supplier<Object> metricsAdjustedForMapping = () -> {
+            var metric = metrics.get();
+
+            var adjusted = new HashMap<String, Number>();
+            for (var metricName : (List<String>) fieldMapping.get("metrics")) {
+                adjusted.put(metricName, metric.get(metricName));
+            }
+
+            return adjusted;
+        };
+
+        if ((Boolean) fieldMapping.getOrDefault("ignore_malformed", false)) {
+            return nullAndMalformedWrapper.apply(metricsAdjustedForMapping).get();
+        }
+
+        return nullWrapper.apply(metricsAdjustedForMapping).get();
+    }
+}