ESQL: Pluggable tests for Operator status

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;
 

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 {
 

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 {
 

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")

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;

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();

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<String> expectedDescriptionOfSimple() {
     protected Matcher<String> expectedToStringOfSimple() {
         return equalTo("ChangePointOperator[channel=0]");
     }
+
+    @Override
+    protected void assertEmptyStatus(Map<String, Object> map) {
+        assertThat(map, nullValue());
+    }
 }

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<String, Object> map) {
+        assertThat(map, nullValue());
+    }
 }

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<String, Object> 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);
     }

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<String, Object> map) {
+        assertMap(map, matchesMap().extraOk().entry("pages_emitted", 0).entry("rows_emitted", 0));
+    }
+}

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<String, Object> map) {
         assumeFalse("not yet standardized", true);
     }
 }

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<InferenceResultsType extends InferenceServiceResults> 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<String, Object> map) {
+        assertMap(map, matchesMap().entry("received_pages", 0).entry("completed_pages", 0).entry("process_nanos", greaterThanOrEqualTo(0)));
     }
 
     @SuppressWarnings("unchecked")