Skip to content

Commit b369cbc

Browse files
idegtiarenkoidegtiarenko
authored
Simplify driver creation in tests (#123457)
1 parent ae46613 commit b369cbc

30 files changed

+231
-298
lines changed

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/Driver.java‎

Lines changed: 0 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -155,37 +155,6 @@ public Driver(
155155
);
156156
}
157157

158-
/**
159-
* Creates a new driver with a chain of operators.
160-
* @param driverContext the driver context
161-
* @param source source operator
162-
* @param intermediateOperators the chain of operators to execute
163-
* @param sink sink operator
164-
* @param releasable a {@link Releasable} to invoked once the chain of operators has run to completion
165-
*/
166-
public Driver(
167-
String taskDescription,
168-
DriverContext driverContext,
169-
SourceOperator source,
170-
List<Operator> intermediateOperators,
171-
SinkOperator sink,
172-
Releasable releasable
173-
) {
174-
this(
175-
"unset",
176-
taskDescription,
177-
System.currentTimeMillis(),
178-
System.nanoTime(),
179-
driverContext,
180-
() -> null,
181-
source,
182-
intermediateOperators,
183-
sink,
184-
DEFAULT_STATUS_INTERVAL,
185-
releasable
186-
);
187-
}
188-
189158
public DriverContext driverContext() {
190159
return driverContext;
191160
}

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@
6363
import org.elasticsearch.compute.test.BlockTestUtils;
6464
import org.elasticsearch.compute.test.OperatorTestCase;
6565
import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
66+
import org.elasticsearch.compute.test.TestDriverFactory;
6667
import org.elasticsearch.core.CheckedConsumer;
6768
import org.elasticsearch.core.Releasables;
6869
import org.elasticsearch.index.mapper.KeywordFieldMapper;
@@ -123,7 +124,7 @@ public void testQueryOperator() throws IOException {
123124
}
124125
});
125126
DriverContext driverContext = driverContext();
126-
drivers.add(new Driver("test", driverContext, factory.get(driverContext), List.of(), docCollector, () -> {}));
127+
drivers.add(TestDriverFactory.create(driverContext, factory.get(driverContext), List.of(), docCollector));
127128
}
128129
OperatorTestCase.runDriver(drivers);
129130
Set<Integer> expectedDocIds = searchForDocIds(reader, query);
@@ -214,8 +215,7 @@ public String toString() {
214215
driverContext
215216
)
216217
);
217-
Driver driver = new Driver(
218-
"test",
218+
Driver driver = TestDriverFactory.create(
219219
driverContext,
220220
luceneOperatorFactory(reader, new MatchAllDocsQuery(), LuceneOperator.NO_LIMIT).get(driverContext),
221221
operators,
@@ -228,8 +228,7 @@ public String toString() {
228228
actualCounts.put(BytesRef.deepCopyOf(spare), counts.getLong(i));
229229
}
230230
page.releaseBlocks();
231-
}),
232-
() -> {}
231+
})
233232
);
234233
OperatorTestCase.runDriver(driver);
235234
assertThat(actualCounts, equalTo(expectedCounts));
@@ -248,8 +247,7 @@ public void testLimitOperator() {
248247
var results = new ArrayList<Long>();
249248
DriverContext driverContext = driverContext();
250249
try (
251-
var driver = new Driver(
252-
"test",
250+
var driver = TestDriverFactory.create(
253251
driverContext,
254252
new SequenceLongBlockSourceOperator(driverContext.blockFactory(), values, 100),
255253
List.of((new LimitOperator.Factory(limit)).get(driverContext)),
@@ -258,8 +256,7 @@ public void testLimitOperator() {
258256
for (int i = 0; i < page.getPositionCount(); i++) {
259257
results.add(block.getLong(i));
260258
}
261-
}),
262-
() -> {}
259+
})
263260
)
264261
) {
265262
OperatorTestCase.runDriver(driver);
@@ -336,8 +333,7 @@ public void testHashLookup() {
336333
var actualValues = new ArrayList<>();
337334
var actualPrimeOrds = new ArrayList<>();
338335
try (
339-
var driver = new Driver(
340-
"test",
336+
var driver = TestDriverFactory.create(
341337
driverContext,
342338
new SequenceLongBlockSourceOperator(driverContext.blockFactory(), values, 100),
343339
List.of(
@@ -354,8 +350,7 @@ public void testHashLookup() {
354350
} finally {
355351
page.releaseBlocks();
356352
}
357-
}),
358-
() -> {}
353+
})
359354
)
360355
) {
361356
OperatorTestCase.runDriver(driver);

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import org.elasticsearch.compute.test.BlockTestUtils;
3232
import org.elasticsearch.compute.test.CannedSourceOperator;
3333
import org.elasticsearch.compute.test.TestBlockFactory;
34+
import org.elasticsearch.compute.test.TestDriverFactory;
3435
import org.elasticsearch.compute.test.TestResultPageSinkOperator;
3536
import org.hamcrest.Matcher;
3637

@@ -110,13 +111,11 @@ public final void testIgnoresNulls() {
110111
List<Page> origInput = BlockTestUtils.deepCopyOf(input, TestBlockFactory.getNonBreakingInstance());
111112

112113
try (
113-
Driver d = new Driver(
114-
"test",
114+
Driver d = TestDriverFactory.create(
115115
driverContext,
116116
new NullInsertingSourceOperator(new CannedSourceOperator(input.iterator()), blockFactory),
117117
List.of(simple().get(driverContext)),
118-
new TestResultPageSinkOperator(results::add),
119-
() -> {}
118+
new TestResultPageSinkOperator(results::add)
120119
)
121120
) {
122121
runDriver(d);

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunctionTests.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.elasticsearch.compute.operator.SequenceIntBlockSourceOperator;
1919
import org.elasticsearch.compute.operator.SourceOperator;
2020
import org.elasticsearch.compute.test.CannedSourceOperator;
21+
import org.elasticsearch.compute.test.TestDriverFactory;
2122

2223
import java.util.List;
2324
import java.util.stream.LongStream;
@@ -64,13 +65,11 @@ public void testRejectsDouble() {
6465
DriverContext driverContext = driverContext();
6566
BlockFactory blockFactory = driverContext.blockFactory();
6667
try (
67-
Driver d = new Driver(
68-
"test",
68+
Driver d = TestDriverFactory.create(
6969
driverContext,
7070
new CannedSourceOperator(Iterators.single(new Page(blockFactory.newDoubleArrayVector(new double[] { 1.0 }, 1).asBlock()))),
7171
List.of(simple().get(driverContext)),
72-
new PageConsumerOperator(page -> fail("shouldn't have made it this far")),
73-
() -> {}
72+
new PageConsumerOperator(page -> fail("shouldn't have made it this far"))
7473
)
7574
) {
7675
expectThrows(Exception.class, () -> runDriver(d)); // ### find a more specific exception type

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunctionTests.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.elasticsearch.compute.operator.SourceOperator;
1919
import org.elasticsearch.compute.test.CannedSourceOperator;
2020
import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
21+
import org.elasticsearch.compute.test.TestDriverFactory;
2122

2223
import java.util.List;
2324
import java.util.stream.LongStream;
@@ -65,13 +66,11 @@ public void testRejectsDouble() {
6566
DriverContext driverContext = driverContext();
6667
BlockFactory blockFactory = driverContext.blockFactory();
6768
try (
68-
Driver d = new Driver(
69-
"test",
69+
Driver d = TestDriverFactory.create(
7070
driverContext,
7171
new CannedSourceOperator(Iterators.single(new Page(blockFactory.newDoubleArrayVector(new double[] { 1.0 }, 1).asBlock()))),
7272
List.of(simple().get(driverContext)),
73-
new PageConsumerOperator(page -> fail("shouldn't have made it this far")),
74-
() -> {}
73+
new PageConsumerOperator(page -> fail("shouldn't have made it this far"))
7574
)
7675
) {
7776
expectThrows(Exception.class, () -> runDriver(d)); // ### find a more specific exception type

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumDoubleAggregatorFunctionTests.java‎

Lines changed: 11 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.elasticsearch.compute.operator.DriverContext;
1616
import org.elasticsearch.compute.operator.SequenceDoubleBlockSourceOperator;
1717
import org.elasticsearch.compute.operator.SourceOperator;
18+
import org.elasticsearch.compute.test.TestDriverFactory;
1819
import org.elasticsearch.compute.test.TestResultPageSinkOperator;
1920
import org.elasticsearch.test.ESTestCase;
2021

@@ -52,13 +53,11 @@ public void testOverflowSucceeds() {
5253
DriverContext driverContext = driverContext();
5354
List<Page> results = new ArrayList<>();
5455
try (
55-
Driver d = new Driver(
56-
"test",
56+
Driver d = TestDriverFactory.create(
5757
driverContext,
5858
new SequenceDoubleBlockSourceOperator(driverContext.blockFactory(), DoubleStream.of(Double.MAX_VALUE - 1, 2)),
5959
List.of(simple().get(driverContext)),
60-
new TestResultPageSinkOperator(results::add),
61-
() -> {}
60+
new TestResultPageSinkOperator(results::add)
6261
)
6362
) {
6463
runDriver(d);
@@ -71,16 +70,14 @@ public void testSummationAccuracy() {
7170
DriverContext driverContext = driverContext();
7271
List<Page> results = new ArrayList<>();
7372
try (
74-
Driver d = new Driver(
75-
"test",
73+
Driver d = TestDriverFactory.create(
7674
driverContext,
7775
new SequenceDoubleBlockSourceOperator(
7876
driverContext.blockFactory(),
7977
DoubleStream.of(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7)
8078
),
8179
List.of(simple().get(driverContext)),
82-
new TestResultPageSinkOperator(results::add),
83-
() -> {}
80+
new TestResultPageSinkOperator(results::add)
8481
)
8582
) {
8683
runDriver(d);
@@ -101,13 +98,11 @@ public void testSummationAccuracy() {
10198
}
10299
driverContext = driverContext();
103100
try (
104-
Driver d = new Driver(
105-
"test",
101+
Driver d = TestDriverFactory.create(
106102
driverContext,
107103
new SequenceDoubleBlockSourceOperator(driverContext.blockFactory(), DoubleStream.of(values)),
108104
List.of(simple().get(driverContext)),
109-
new TestResultPageSinkOperator(results::add),
110-
() -> {}
105+
new TestResultPageSinkOperator(results::add)
111106
)
112107
) {
113108
runDriver(d);
@@ -124,13 +119,11 @@ public void testSummationAccuracy() {
124119
}
125120
driverContext = driverContext();
126121
try (
127-
Driver d = new Driver(
128-
"test",
122+
Driver d = TestDriverFactory.create(
129123
driverContext,
130124
new SequenceDoubleBlockSourceOperator(driverContext.blockFactory(), DoubleStream.of(largeValues)),
131125
List.of(simple().get(driverContext)),
132-
new TestResultPageSinkOperator(results::add),
133-
() -> {}
126+
new TestResultPageSinkOperator(results::add)
134127
)
135128
) {
136129
runDriver(d);
@@ -144,13 +137,11 @@ public void testSummationAccuracy() {
144137
}
145138
driverContext = driverContext();
146139
try (
147-
Driver d = new Driver(
148-
"test",
140+
Driver d = TestDriverFactory.create(
149141
driverContext,
150142
new SequenceDoubleBlockSourceOperator(driverContext.blockFactory(), DoubleStream.of(largeValues)),
151143
List.of(simple().get(driverContext)),
152-
new TestResultPageSinkOperator(results::add),
153-
() -> {}
144+
new TestResultPageSinkOperator(results::add)
154145
)
155146
) {
156147
runDriver(d);

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumFloatAggregatorFunctionTests.java‎

Lines changed: 11 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.elasticsearch.compute.operator.DriverContext;
1616
import org.elasticsearch.compute.operator.SequenceFloatBlockSourceOperator;
1717
import org.elasticsearch.compute.operator.SourceOperator;
18+
import org.elasticsearch.compute.test.TestDriverFactory;
1819
import org.elasticsearch.compute.test.TestResultPageSinkOperator;
1920
import org.elasticsearch.test.ESTestCase;
2021

@@ -52,13 +53,11 @@ public void testOverflowSucceeds() {
5253
DriverContext driverContext = driverContext();
5354
List<Page> results = new ArrayList<>();
5455
try (
55-
Driver d = new Driver(
56-
"test",
56+
Driver d = TestDriverFactory.create(
5757
driverContext,
5858
new SequenceFloatBlockSourceOperator(driverContext.blockFactory(), Stream.of(Float.MAX_VALUE - 1, 2f)),
5959
List.of(simple().get(driverContext)),
60-
new TestResultPageSinkOperator(results::add),
61-
() -> {}
60+
new TestResultPageSinkOperator(results::add)
6261
)
6362
) {
6463
runDriver(d);
@@ -71,16 +70,14 @@ public void testSummationAccuracy() {
7170
DriverContext driverContext = driverContext();
7271
List<Page> results = new ArrayList<>();
7372
try (
74-
Driver d = new Driver(
75-
"test",
73+
Driver d = TestDriverFactory.create(
7674
driverContext,
7775
new SequenceFloatBlockSourceOperator(
7876
driverContext.blockFactory(),
7977
Stream.of(0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f, 1.3f, 1.4f, 1.5f, 1.6f, 1.7f)
8078
),
8179
List.of(simple().get(driverContext)),
82-
new TestResultPageSinkOperator(results::add),
83-
() -> {}
80+
new TestResultPageSinkOperator(results::add)
8481
)
8582
) {
8683
runDriver(d);
@@ -101,13 +98,11 @@ public void testSummationAccuracy() {
10198
}
10299
driverContext = driverContext();
103100
try (
104-
Driver d = new Driver(
105-
"test",
101+
Driver d = TestDriverFactory.create(
106102
driverContext,
107103
new SequenceFloatBlockSourceOperator(driverContext.blockFactory(), Stream.of(values)),
108104
List.of(simple().get(driverContext)),
109-
new TestResultPageSinkOperator(results::add),
110-
() -> {}
105+
new TestResultPageSinkOperator(results::add)
111106
)
112107
) {
113108
runDriver(d);
@@ -124,13 +119,11 @@ public void testSummationAccuracy() {
124119
}
125120
driverContext = driverContext();
126121
try (
127-
Driver d = new Driver(
128-
"test",
122+
Driver d = TestDriverFactory.create(
129123
driverContext,
130124
new SequenceFloatBlockSourceOperator(driverContext.blockFactory(), Stream.of(largeValues)),
131125
List.of(simple().get(driverContext)),
132-
new TestResultPageSinkOperator(results::add),
133-
() -> {}
126+
new TestResultPageSinkOperator(results::add)
134127
)
135128
) {
136129
runDriver(d);
@@ -144,13 +137,11 @@ public void testSummationAccuracy() {
144137
}
145138
driverContext = driverContext();
146139
try (
147-
Driver d = new Driver(
148-
"test",
140+
Driver d = TestDriverFactory.create(
149141
driverContext,
150142
new SequenceFloatBlockSourceOperator(driverContext.blockFactory(), Stream.of(largeValues)),
151143
List.of(simple().get(driverContext)),
152-
new TestResultPageSinkOperator(results::add),
153-
() -> {}
144+
new TestResultPageSinkOperator(results::add)
154145
)
155146
) {
156147
runDriver(d);

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumIntAggregatorFunctionTests.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.elasticsearch.compute.operator.SequenceIntBlockSourceOperator;
1919
import org.elasticsearch.compute.operator.SourceOperator;
2020
import org.elasticsearch.compute.test.CannedSourceOperator;
21+
import org.elasticsearch.compute.test.TestDriverFactory;
2122

2223
import java.util.List;
2324
import java.util.stream.LongStream;
@@ -51,13 +52,11 @@ public void testRejectsDouble() {
5152
DriverContext driverContext = driverContext();
5253
BlockFactory blockFactory = driverContext.blockFactory();
5354
try (
54-
Driver d = new Driver(
55-
"test",
55+
Driver d = TestDriverFactory.create(
5656
driverContext,
5757
new CannedSourceOperator(Iterators.single(new Page(blockFactory.newDoubleArrayVector(new double[] { 1.0 }, 1).asBlock()))),
5858
List.of(simple().get(driverContext)),
59-
new PageConsumerOperator(page -> fail("shouldn't have made it this far")),
60-
() -> {}
59+
new PageConsumerOperator(page -> fail("shouldn't have made it this far"))
6160
)
6261
) {
6362
expectThrows(Exception.class, () -> runDriver(d)); // ### find a more specific exception type

0 commit comments

Comments
 (0)