ESQL: Enable block tracking in extract tests

This enables memory tracking in the tests for `grok` and `dissect`. They are already tracking their memory.
nik9000 · Oct 26, 2023 · 5a6a433 · 5a6a433
1 parent 328ebc4
commit 5a6a433
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 26 deletions.
diff --git a/.../compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java b/.../compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java
@@ -91,4 +91,9 @@ protected void assertSimpleOutput(List<Page> input, List<Page> results) {
     protected ByteSizeValue smallEnoughToCircuitBreak() {
         return ByteSizeValue.ofBytes(between(1, 32));
     }
+
+    @Override
+    protected DriverContext driverContext() {
+        return breakingDriverContext();
+    }
 }
diff --git a/.../compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java b/.../compute/src/test/java/org/elasticsearch/compute/operator/StringExtractOperatorTests.java
@@ -99,31 +99,40 @@ public Block.Ref eval(Page page) {
             public void close() {}
         }, new FirstWord("test"), driverContext());
 
-        BytesRefBlock.Builder builder = BytesRefBlock.newBlockBuilder(1);
-        builder.beginPositionEntry();
-        builder.appendBytesRef(new BytesRef("foo1 bar1"));
-        builder.appendBytesRef(new BytesRef("foo2 bar2"));
-        builder.endPositionEntry();
-        builder.beginPositionEntry();
-        builder.appendBytesRef(new BytesRef("foo3 bar3"));
-        builder.appendBytesRef(new BytesRef("foo4 bar4"));
-        builder.appendBytesRef(new BytesRef("foo5 bar5"));
-        builder.endPositionEntry();
-        Page page = new Page(builder.build());
-
-        Page result = operator.process(page);
-        Block resultBlock = result.getBlock(1);
-        assertThat(resultBlock.getPositionCount(), equalTo(2));
-        assertThat(resultBlock.getValueCount(0), equalTo(2));
-        assertThat(resultBlock.getValueCount(1), equalTo(3));
-        BytesRefBlock brb = (BytesRefBlock) resultBlock;
-        BytesRef spare = new BytesRef("");
-        int idx = brb.getFirstValueIndex(0);
-        assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo1"));
-        assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo2"));
-        idx = brb.getFirstValueIndex(1);
-        assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo3"));
-        assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo4"));
-        assertThat(brb.getBytesRef(idx + 2, spare).utf8ToString(), equalTo("foo5"));
+        Page result = null;
+        try (BytesRefBlock.Builder builder = BytesRefBlock.newBlockBuilder(1)) {
+            builder.beginPositionEntry();
+            builder.appendBytesRef(new BytesRef("foo1 bar1"));
+            builder.appendBytesRef(new BytesRef("foo2 bar2"));
+            builder.endPositionEntry();
+            builder.beginPositionEntry();
+            builder.appendBytesRef(new BytesRef("foo3 bar3"));
+            builder.appendBytesRef(new BytesRef("foo4 bar4"));
+            builder.appendBytesRef(new BytesRef("foo5 bar5"));
+            builder.endPositionEntry();
+            result = operator.process(new Page(builder.build()));
+        }
+        try {
+            Block resultBlock = result.getBlock(1);
+            assertThat(resultBlock.getPositionCount(), equalTo(2));
+            assertThat(resultBlock.getValueCount(0), equalTo(2));
+            assertThat(resultBlock.getValueCount(1), equalTo(3));
+            BytesRefBlock brb = (BytesRefBlock) resultBlock;
+            BytesRef spare = new BytesRef("");
+            int idx = brb.getFirstValueIndex(0);
+            assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo1"));
+            assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo2"));
+            idx = brb.getFirstValueIndex(1);
+            assertThat(brb.getBytesRef(idx, spare).utf8ToString(), equalTo("foo3"));
+            assertThat(brb.getBytesRef(idx + 1, spare).utf8ToString(), equalTo("foo4"));
+            assertThat(brb.getBytesRef(idx + 2, spare).utf8ToString(), equalTo("foo5"));
+        } finally {
+            result.releaseBlocks();
+        }
+    }
+
+    @Override
+    protected DriverContext driverContext() {
+        return breakingDriverContext();
     }
 }