diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneCountOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneCountOperatorTests.java index 2487016004cf2..de5de4743f971 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneCountOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneCountOperatorTests.java @@ -21,8 +21,8 @@ import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.Driver; import org.elasticsearch.compute.operator.DriverContext; -import org.elasticsearch.compute.test.AnyOperatorTestCase; import org.elasticsearch.compute.test.OperatorTestCase; +import org.elasticsearch.compute.test.SourceOperatorTestCase; import org.elasticsearch.compute.test.TestDriverFactory; import org.elasticsearch.compute.test.TestResultPageSinkOperator; import org.elasticsearch.core.IOUtils; @@ -41,7 +41,7 @@ import static org.hamcrest.Matchers.lessThanOrEqualTo; import static org.hamcrest.Matchers.matchesRegex; -public class LuceneCountOperatorTests extends AnyOperatorTestCase { +public class LuceneCountOperatorTests extends SourceOperatorTestCase { private Directory directory = newDirectory(); private IndexReader reader; diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMaxOperatorTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMaxOperatorTestCase.java index ea717ae805118..aac796d699239 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMaxOperatorTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMaxOperatorTestCase.java @@ -22,8 +22,8 @@ import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.Driver; import org.elasticsearch.compute.operator.DriverContext; -import org.elasticsearch.compute.test.AnyOperatorTestCase; import org.elasticsearch.compute.test.OperatorTestCase; +import org.elasticsearch.compute.test.SourceOperatorTestCase; import org.elasticsearch.compute.test.TestDriverFactory; import org.elasticsearch.compute.test.TestResultPageSinkOperator; import org.elasticsearch.core.IOUtils; @@ -43,7 +43,7 @@ import static org.hamcrest.Matchers.lessThanOrEqualTo; import static org.hamcrest.Matchers.matchesRegex; -public abstract class LuceneMaxOperatorTestCase extends AnyOperatorTestCase { +public abstract class LuceneMaxOperatorTestCase extends SourceOperatorTestCase { protected interface NumberTypeTest { diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMinOperatorTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMinOperatorTestCase.java index ac0d16103ed92..a19649bb2c8bf 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMinOperatorTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneMinOperatorTestCase.java @@ -22,8 +22,8 @@ import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.Driver; import org.elasticsearch.compute.operator.DriverContext; -import org.elasticsearch.compute.test.AnyOperatorTestCase; import org.elasticsearch.compute.test.OperatorTestCase; +import org.elasticsearch.compute.test.SourceOperatorTestCase; import org.elasticsearch.compute.test.TestDriverFactory; import org.elasticsearch.compute.test.TestResultPageSinkOperator; import org.elasticsearch.core.IOUtils; @@ -43,7 +43,7 @@ import static org.hamcrest.Matchers.lessThanOrEqualTo; import static org.hamcrest.Matchers.matchesRegex; -public abstract class LuceneMinOperatorTestCase extends AnyOperatorTestCase { +public abstract class LuceneMinOperatorTestCase extends SourceOperatorTestCase { protected interface NumberTypeTest { diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java index cc21d180113a0..db34c2c960c59 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java @@ -31,8 +31,8 @@ import org.elasticsearch.compute.operator.PageConsumerOperator; import org.elasticsearch.compute.operator.SinkOperator; import org.elasticsearch.compute.operator.SourceOperator; -import org.elasticsearch.compute.test.AnyOperatorTestCase; import org.elasticsearch.compute.test.OperatorTestCase; +import org.elasticsearch.compute.test.SourceOperatorTestCase; import org.elasticsearch.compute.test.TestDriverFactory; import org.elasticsearch.compute.test.TestResultPageSinkOperator; import org.elasticsearch.core.IOUtils; @@ -64,7 +64,7 @@ import static org.hamcrest.Matchers.lessThanOrEqualTo; import static org.hamcrest.Matchers.matchesRegex; -public class LuceneSourceOperatorTests extends AnyOperatorTestCase { +public class LuceneSourceOperatorTests extends SourceOperatorTestCase { private static final MappedFieldType S_FIELD = new NumberFieldMapper.NumberFieldType("s", NumberFieldMapper.NumberType.LONG); @ParametersFactory(argumentFormatting = "%s %s") diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java index 26540caee0b1f..414bf82fcbcf6 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java @@ -27,8 +27,8 @@ import org.elasticsearch.compute.lucene.read.ValuesSourceReaderOperatorTests; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.Operator; -import org.elasticsearch.compute.test.AnyOperatorTestCase; import org.elasticsearch.compute.test.OperatorTestCase; +import org.elasticsearch.compute.test.SourceOperatorTestCase; import org.elasticsearch.compute.test.TestDriverFactory; import org.elasticsearch.compute.test.TestResultPageSinkOperator; import org.elasticsearch.core.IOUtils; @@ -52,7 +52,7 @@ import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.matchesRegex; -public class LuceneTopNSourceOperatorTests extends AnyOperatorTestCase { +public class LuceneTopNSourceOperatorTests extends SourceOperatorTestCase { private static final MappedFieldType S_FIELD = new NumberFieldMapper.NumberFieldType("s", NumberFieldMapper.NumberType.LONG); private Directory directory = newDirectory(); private IndexReader reader; diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorTests.java index 3a3fac1afc595..ea624288a8c82 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorTests.java @@ -41,8 +41,8 @@ import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.DriverStatus; import org.elasticsearch.compute.operator.Operator; -import org.elasticsearch.compute.test.AnyOperatorTestCase; import org.elasticsearch.compute.test.OperatorTestCase; +import org.elasticsearch.compute.test.SourceOperatorTestCase; import org.elasticsearch.compute.test.TestDriverFactory; import org.elasticsearch.compute.test.TestResultPageSinkOperator; import org.elasticsearch.core.CheckedFunction; @@ -74,7 +74,7 @@ import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.lessThanOrEqualTo; -public class TimeSeriesSourceOperatorTests extends AnyOperatorTestCase { +public class TimeSeriesSourceOperatorTests extends SourceOperatorTestCase { private IndexReader reader; private final Directory directory = newDirectory(); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ChangePointOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ChangePointOperatorTests.java index d3e68fc603cf9..c383028900199 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ChangePointOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ChangePointOperatorTests.java @@ -19,9 +19,11 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.hasSize; +import static org.hamcrest.Matchers.nullValue; public class ChangePointOperatorTests extends OperatorTestCase { @@ -83,4 +85,9 @@ protected Matcher expectedDescriptionOfSimple() { protected Matcher expectedToStringOfSimple() { return equalTo("ChangePointOperator[channel=0]"); } + + @Override + protected void assertEmptyStatus(Map map) { + assertThat(map, nullValue()); + } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/OutputOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/OutputOperatorTests.java index 8947f7ca5aded..3116b247020f4 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/OutputOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/OutputOperatorTests.java @@ -11,9 +11,11 @@ import org.hamcrest.Matcher; import java.util.List; +import java.util.Map; import java.util.stream.IntStream; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.nullValue; public class OutputOperatorTests extends AnyOperatorTestCase { @Override @@ -48,4 +50,9 @@ public void testBigToString() { public void testBigDescription() { assertThat(big().describe(), equalTo(expectedDescriptionOfBig())); } + + @Override + protected void assertEmptyStatus(Map map) { + assertThat(map, nullValue()); + } } diff --git a/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/AnyOperatorTestCase.java b/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/AnyOperatorTestCase.java index 04c2323116ab7..d047a01ec7db6 100644 --- a/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/AnyOperatorTestCase.java +++ b/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/AnyOperatorTestCase.java @@ -14,18 +14,19 @@ import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.Operator; -import org.elasticsearch.compute.operator.SinkOperator; import org.elasticsearch.compute.operator.SourceOperator; +import org.elasticsearch.core.Nullable; +import org.elasticsearch.test.MapMatcher; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentType; import org.hamcrest.Matcher; import java.io.IOException; +import java.util.Map; -import static org.hamcrest.Matchers.both; -import static org.hamcrest.Matchers.either; -import static org.hamcrest.Matchers.hasKey; +import static org.elasticsearch.test.MapMatcher.assertMap; +import static org.elasticsearch.test.MapMatcher.matchesMap; import static org.hamcrest.Matchers.matchesPattern; /** @@ -107,39 +108,44 @@ public final void testSimpleToString() { /** * Ensures that the Operator.Status of this operator has the standard fields. */ - public void testOperatorStatus() throws IOException { + public final void testOperatorStatus() throws IOException { DriverContext driverContext = driverContext(); try (var operator = simple().get(driverContext)) { Operator.Status status = operator.status(); + if (status == null) { + assertEmptyStatus(null); + return; + } - assumeTrue("Operator does not provide a status", status != null); - - var xContent = XContentType.JSON.xContent(); - try (var xContentBuilder = XContentBuilder.builder(xContent)) { + try (var xContentBuilder = XContentBuilder.builder(XContentType.JSON.xContent())) { status.toXContent(xContentBuilder, ToXContent.EMPTY_PARAMS); var bytesReference = BytesReference.bytes(xContentBuilder); var map = XContentHelper.convertToMap(bytesReference, false, xContentBuilder.contentType()).v2(); - if (operator instanceof SourceOperator) { - assertThat(map, hasKey("pages_emitted")); - assertThat(map, hasKey("rows_emitted")); - } else if (operator instanceof SinkOperator) { - assertThat(map, hasKey("pages_received")); - assertThat(map, hasKey("rows_received")); - } else { - assertThat(map, either(hasKey("pages_processed")).or(both(hasKey("pages_received")).and(hasKey("pages_emitted")))); - assertThat(map, hasKey("rows_received")); - assertThat(map, hasKey("rows_emitted")); - } + assertEmptyStatus(map); } } } + /** + * Assert that the status is sane. + */ + protected void assertEmptyStatus(@Nullable Map map) { + MapMatcher matcher = matchesMap().extraOk(); + if (map.containsKey("pages_processed")) { + matcher = matcher.entry("pages_processed", 0); + } else { + matcher = matcher.entry("pages_received", 0).entry("pages_emitted", 0); + } + matcher = matcher.entry("rows_received", 0).entry("rows_emitted", 0); + assertMap(map, matcher); + } + /** * A {@link DriverContext} with a nonBreakingBigArrays. */ - protected DriverContext driverContext() { // TODO make this final once all operators support memory tracking + protected final DriverContext driverContext() { BlockFactory blockFactory = blockFactory(); return new DriverContext(blockFactory.bigArrays(), blockFactory); } diff --git a/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/SourceOperatorTestCase.java b/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/SourceOperatorTestCase.java new file mode 100644 index 0000000000000..904aa8882c857 --- /dev/null +++ b/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/SourceOperatorTestCase.java @@ -0,0 +1,20 @@ +/* + * 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.compute.test; + +import java.util.Map; + +import static org.elasticsearch.test.MapMatcher.assertMap; +import static org.elasticsearch.test.MapMatcher.matchesMap; + +public abstract class SourceOperatorTestCase extends AnyOperatorTestCase { + @Override + protected void assertEmptyStatus(Map map) { + assertMap(map, matchesMap().extraOk().entry("pages_emitted", 0).entry("rows_emitted", 0)); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperatorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperatorTests.java index c4a381eccff65..43e6e914d03a3 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperatorTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperatorTests.java @@ -304,7 +304,7 @@ public void release() { } @Override - public void testOperatorStatus() { + protected void assertEmptyStatus(Map map) { assumeFalse("not yet standardized", true); } } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java index 6ec951a617a62..a95935c1e22e7 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java @@ -25,8 +25,6 @@ import org.elasticsearch.compute.data.IntBlock; import org.elasticsearch.compute.data.LongBlock; import org.elasticsearch.compute.data.Page; -import org.elasticsearch.compute.operator.AsyncOperator; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.compute.operator.SourceOperator; import org.elasticsearch.compute.test.AbstractBlockSourceOperator; @@ -43,12 +41,14 @@ import org.junit.After; import org.junit.Before; +import java.util.Map; import java.util.function.BiFunction; import java.util.function.Consumer; +import static org.elasticsearch.test.MapMatcher.assertMap; +import static org.elasticsearch.test.MapMatcher.matchesMap; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThanOrEqualTo; -import static org.hamcrest.Matchers.notNullValue; public abstract class InferenceOperatorTestCase extends OperatorTestCase { protected ThreadPool threadPool; @@ -121,16 +121,8 @@ protected Page createPage(int positionOffset, int length) { } @Override - public void testOperatorStatus() { - DriverContext driverContext = driverContext(); - try (var operator = simple().get(driverContext)) { - AsyncOperator.Status status = asInstanceOf(AsyncOperator.Status.class, operator.status()); - - assertThat(status, notNullValue()); - assertThat(status.receivedPages(), equalTo(0L)); - assertThat(status.completedPages(), equalTo(0L)); - assertThat(status.procesNanos(), greaterThanOrEqualTo(0L)); - } + protected void assertEmptyStatus(Map map) { + assertMap(map, matchesMap().entry("received_pages", 0).entry("completed_pages", 0).entry("process_nanos", greaterThanOrEqualTo(0))); } @SuppressWarnings("unchecked")