Skip to content

Commit cbe757b

Browse files
leontyevdvleontyevdv
committed
ES|QL: Refactor tests to
- Use TypeSignature instead of Map.Entry - Refactor LogicalPlanOptimizerTests to use Capabilities
1 parent 03f45d9 commit cbe757b

File tree

8 files changed

+74
-86
lines changed

8 files changed

+74
-86
lines changed

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/LookupJoinTypesIT.java‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import java.util.HashSet;
3434
import java.util.Iterator;
3535
import java.util.LinkedHashMap;
36+
import java.util.LinkedHashSet;
3637
import java.util.List;
3738
import java.util.Locale;
3839
import java.util.Map;
@@ -270,16 +271,21 @@ private static boolean existingIndex(Collection<TestConfigs> existing, DataType
270271

271272
/** This test generates documentation for the supported output types of the lookup join. */
272273
public void testOutputSupportedTypes() throws Exception {
273-
Map<List<DocsV3Support.Param>, DataType> signatures = new LinkedHashMap<>();
274+
Set<DocsV3Support.TypeSignature> signatures = new LinkedHashSet<>();
274275
for (TestConfigs configs : testConfigurations.values()) {
275276
if (configs.group.equals("unsupported") || configs.group.equals("union-types")) {
276277
continue;
277278
}
278279
for (TestConfig config : configs.configs.values()) {
279280
if (config instanceof TestConfigPasses) {
280-
signatures.put(
281-
List.of(new DocsV3Support.Param(config.mainType(), List.of()), new DocsV3Support.Param(config.lookupType(), null)),
282-
null
281+
signatures.add(
282+
new DocsV3Support.TypeSignature(
283+
List.of(
284+
new DocsV3Support.Param(config.mainType(), List.of()),
285+
new DocsV3Support.Param(config.lookupType(), null)
286+
),
287+
null
288+
)
283289
);
284290
}
285291
}
@@ -770,7 +776,7 @@ private boolean isValidDataType(DataType dataType) {
770776
return UNDER_CONSTRUCTION.get(dataType) == null || UNDER_CONSTRUCTION.get(dataType).isEnabled();
771777
}
772778

773-
private static void saveJoinTypes(Supplier<Map<List<DocsV3Support.Param>, DataType>> signatures) throws Exception {
779+
private static void saveJoinTypes(Supplier<Set<DocsV3Support.TypeSignature>> signatures) throws Exception {
774780
if (System.getProperty("generateDocs") == null) {
775781
return;
776782
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,6 @@
6666
import java.util.ArrayList;
6767
import java.util.Arrays;
6868
import java.util.Collections;
69-
import java.util.HashMap;
7069
import java.util.HashSet;
7170
import java.util.List;
7271
import java.util.Locale;
@@ -758,13 +757,13 @@ public static void testFunctionInfo() {
758757
for (int i = 0; i < args.size(); i++) {
759758
typesFromSignature.add(new HashSet<>());
760759
}
761-
for (Map.Entry<List<DocsV3Support.Param>, DataType> entry : signatures(testClass).entrySet()) {
762-
List<DocsV3Support.Param> types = entry.getKey();
760+
for (DocsV3Support.TypeSignature entry : signatures(testClass)) {
761+
List<DocsV3Support.Param> types = entry.argTypes();
763762
for (int i = 0; i < args.size() && i < types.size(); i++) {
764763
typesFromSignature.get(i).add(types.get(i).dataType().esNameIfPossible());
765764
}
766-
if (DataType.UNDER_CONSTRUCTION.containsKey(entry.getValue()) == false) {
767-
returnFromSignature.add(entry.getValue().esNameIfPossible());
765+
if (DataType.UNDER_CONSTRUCTION.containsKey(entry.returnType()) == false) {
766+
returnFromSignature.add(entry.returnType().esNameIfPossible());
768767
}
769768
}
770769

@@ -838,8 +837,9 @@ public static void testFunctionLicenseChecks() throws Exception {
838837
TestCheckLicense checkLicense = new TestCheckLicense();
839838

840839
// Go through all signatures and assert that the license is as expected
841-
signatures(testClass).forEach((signature, returnType) -> {
840+
signatures(testClass).forEach((signatureItem) -> {
842841
try {
842+
List<DocsV3Support.Param> signature = signatureItem.argTypes();
843843
License.OperationMode license = licenseChecker.invoke(signature.stream().map(DocsV3Support.Param::dataType).toList());
844844
assertNotNull("License should not be null", license);
845845

@@ -933,14 +933,14 @@ protected final void assertTypeResolutionFailure(Expression expression) {
933933
/**
934934
* Unique signatures in this test’s parameters.
935935
*/
936-
private static Map<List<DocsV3Support.Param>, DataType> signatures;
936+
private static Set<DocsV3Support.TypeSignature> signatures;
937937

938-
public static Map<List<DocsV3Support.Param>, DataType> signatures(Class<?> testClass) {
938+
public static Set<DocsV3Support.TypeSignature> signatures(Class<?> testClass) {
939939
if (signatures != null && classGeneratingSignatures == testClass) {
940940
return signatures;
941941
}
942942
classGeneratingSignatures = testClass;
943-
signatures = new HashMap<>();
943+
signatures = new HashSet<>();
944944
Set<Method> paramsFactories = new ClassModel(testClass).getAnnotatedLeafMethods(ParametersFactory.class).keySet();
945945
assertThat(paramsFactories, hasSize(1));
946946
Method paramsFactory = paramsFactories.iterator().next();
@@ -960,7 +960,7 @@ public static Map<List<DocsV3Support.Param>, DataType> signatures(Class<?> testC
960960
continue;
961961
}
962962
List<DocsV3Support.Param> sig = tc.getData().stream().map(d -> new DocsV3Support.Param(d.type(), d.appliesTo())).toList();
963-
signatures.putIfAbsent(signatureTypes(testClass, sig), tc.expectedType());
963+
signatures.add(new DocsV3Support.TypeSignature(signatureTypes(testClass, sig), tc.expectedType()));
964964
}
965965
return signatures;
966966
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/DocsV3Support.java‎

Lines changed: 40 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
import java.util.Map;
6565
import java.util.Objects;
6666
import java.util.Optional;
67+
import java.util.Set;
6768
import java.util.TreeMap;
6869
import java.util.function.Function;
6970
import java.util.function.Supplier;
@@ -101,6 +102,8 @@
101102
public abstract class DocsV3Support {
102103
public record Param(DataType dataType, List<FunctionAppliesTo> appliesTo) {}
103104

105+
public record TypeSignature(List<DocsV3Support.Param> argTypes, DataType returnType) {}
106+
104107
private static final Logger logger = LogManager.getLogger(DocsV3Support.class);
105108

106109
private static final String DOCS_WARNING_JSON =
@@ -374,15 +377,15 @@ public License.OperationMode invoke(List<DataType> fieldTypes) throws Exception
374377
protected final String category;
375378
protected final String name;
376379
protected final FunctionDefinition definition;
377-
protected final Supplier<Map<List<Param>, DataType>> signatures;
380+
protected final Supplier<Set<TypeSignature>> signatures;
378381
protected final Callbacks callbacks;
379382
private final LicenseRequirementChecker licenseChecker;
380383

381384
protected DocsV3Support(
382385
String category,
383386
String name,
384387
Class<?> testClass,
385-
Supplier<Map<List<Param>, DataType>> signatures,
388+
Supplier<Set<TypeSignature>> signatures,
386389
Callbacks callbacks
387390
) {
388391
this(category, name, null, testClass, signatures, callbacks);
@@ -393,7 +396,7 @@ private DocsV3Support(
393396
String name,
394397
FunctionDefinition definition,
395398
Class<?> testClass,
396-
Supplier<Map<List<Param>, DataType>> signatures,
399+
Supplier<Set<TypeSignature>> signatures,
397400
Callbacks callbacks
398401
) {
399402
this.category = category;
@@ -573,7 +576,7 @@ private FunctionDocsSupport(String name, Class<?> testClass, Callbacks callbacks
573576
String name,
574577
Class<?> testClass,
575578
FunctionDefinition definition,
576-
Supplier<Map<List<Param>, DataType>> signatures,
579+
Supplier<Set<TypeSignature>> signatures,
577580
Callbacks callbacks
578581
) {
579582
super("functions", name, definition, testClass, signatures, callbacks);
@@ -780,7 +783,7 @@ public OperatorsDocsSupport(
780783
String name,
781784
Class<?> testClass,
782785
OperatorConfig op,
783-
Supplier<Map<List<Param>, DataType>> signatures,
786+
Supplier<Set<TypeSignature>> signatures,
784787
Callbacks callbacks
785788
) {
786789
super("operators", name, testClass, signatures, callbacks);
@@ -971,7 +974,7 @@ public CommandsDocsSupport(
971974
ObservabilityTier observabilityTier,
972975
Callbacks callbacks
973976
) {
974-
super("commands", name, testClass, Map::of, callbacks);
977+
super("commands", name, testClass, Set::of, callbacks);
975978
this.command = command;
976979
this.licenseState = licenseState;
977980
this.observabilityTier = observabilityTier;
@@ -983,7 +986,7 @@ public CommandsDocsSupport(
983986
Class<?> testClass,
984987
LogicalPlan command,
985988
List<EsqlFunctionRegistry.ArgSignature> args,
986-
Supplier<Map<List<Param>, DataType>> signatures,
989+
Supplier<Set<TypeSignature>> signatures,
987990
Callbacks callbacks
988991
) {
989992
super("commands", name, testClass, signatures, callbacks);
@@ -1043,12 +1046,12 @@ void renderTypes(String name, List<EsqlFunctionRegistry.ArgSignature> args) thro
10431046
}
10441047

10451048
Map<String, List<String>> compactedTable = new TreeMap<>();
1046-
for (Map.Entry<List<DocsV3Support.Param>, DataType> sig : this.signatures.get().entrySet()) {
1047-
if (shouldHideSignature(sig.getKey(), sig.getValue())) {
1049+
for (TypeSignature sig : this.signatures.get()) {
1050+
if (shouldHideSignature(sig.argTypes(), sig.returnType())) {
10481051
continue;
10491052
}
1050-
String mainType = sig.getKey().getFirst().dataType().esNameIfPossible();
1051-
String secondaryType = sig.getKey().get(1).dataType().esNameIfPossible();
1053+
String mainType = sig.argTypes().getFirst().dataType().esNameIfPossible();
1054+
String secondaryType = sig.argTypes().get(1).dataType().esNameIfPossible();
10521055
List<String> secondaryTypes = compactedTable.computeIfAbsent(mainType, (k) -> new ArrayList<>());
10531056
secondaryTypes.add(secondaryType);
10541057
}
@@ -1092,7 +1095,7 @@ void renderParametersList(List<String> argNames, List<String> argDescriptions) t
10921095
}
10931096

10941097
void renderTypes(String name, List<EsqlFunctionRegistry.ArgSignature> args) throws IOException {
1095-
boolean showResultColumn = signatures.get().values().stream().anyMatch(Objects::nonNull);
1098+
boolean showResultColumn = signatures.get().stream().map(TypeSignature::returnType).anyMatch(Objects::nonNull);
10961099
StringBuilder header = new StringBuilder("| ");
10971100
StringBuilder separator = new StringBuilder("| ");
10981101
List<String> argNames = args.stream().map(EsqlFunctionRegistry.ArgSignature::name).toList();
@@ -1106,11 +1109,11 @@ void renderTypes(String name, List<EsqlFunctionRegistry.ArgSignature> args) thro
11061109
}
11071110

11081111
List<String> table = new ArrayList<>();
1109-
for (Map.Entry<List<DocsV3Support.Param>, DataType> sig : this.signatures.get().entrySet()) { // TODO flip to using sortedSignatures
1110-
if (shouldHideSignature(sig.getKey(), sig.getValue())) {
1112+
for (TypeSignature sig : this.signatures.get()) { // TODO flip to using sortedSignatures
1113+
if (shouldHideSignature(sig.argTypes(), sig.returnType())) {
11111114
continue;
11121115
}
1113-
if (sig.getKey().size() > argNames.size()) { // skip variadic [test] cases (but not those with optional parameters)
1116+
if (sig.argTypes().size() > argNames.size()) { // skip variadic [test] cases (but not those with optional parameters)
11141117
continue;
11151118
}
11161119
table.add(getTypeRow(args, sig, argNames, showResultColumn));
@@ -1131,13 +1134,13 @@ void renderTypes(String name, List<EsqlFunctionRegistry.ArgSignature> args) thro
11311134

11321135
private static String getTypeRow(
11331136
List<EsqlFunctionRegistry.ArgSignature> args,
1134-
Map.Entry<List<Param>, DataType> sig,
1137+
TypeSignature sig,
11351138
List<String> argNames,
11361139
boolean showResultColumn
11371140
) {
11381141
StringBuilder b = new StringBuilder("| ");
1139-
for (int i = 0; i < sig.getKey().size(); i++) {
1140-
Param param = sig.getKey().get(i);
1142+
for (int i = 0; i < sig.argTypes().size(); i++) {
1143+
Param param = sig.argTypes().get(i);
11411144
EsqlFunctionRegistry.ArgSignature argSignature = args.get(i);
11421145
if (argSignature.mapArg()) {
11431146
b.append("named parameters");
@@ -1149,9 +1152,9 @@ private static String getTypeRow(
11491152
}
11501153
b.append(" | ");
11511154
}
1152-
b.append("| ".repeat(argNames.size() - sig.getKey().size()));
1155+
b.append("| ".repeat(argNames.size() - sig.argTypes().size()));
11531156
if (showResultColumn) {
1154-
b.append(sig.getValue().esNameIfPossible());
1157+
b.append(sig.returnType().esNameIfPossible());
11551158
b.append(" |");
11561159
}
11571160
return b.toString();
@@ -1300,24 +1303,24 @@ void renderKibanaFunctionDefinition(
13001303
builder.startArray("params");
13011304
builder.endArray();
13021305
// There should only be one return type so just use that as the example
1303-
builder.field("returnType", signatures.get().values().iterator().next().esNameIfPossible());
1306+
builder.field("returnType", signatures.get().iterator().next().returnType().esNameIfPossible());
13041307
builder.endObject();
13051308
} else {
13061309
int minArgCount = (int) args.stream().filter(a -> false == a.optional()).count();
1307-
for (Map.Entry<List<DocsV3Support.Param>, DataType> sig : sortedSignatures()) {
1308-
if (variadic && sig.getKey().size() > args.size()) {
1310+
for (TypeSignature sig : sortedSignatures()) {
1311+
if (variadic && sig.argTypes().size() > args.size()) {
13091312
// For variadic functions we test much longer signatures, let’s just stop at the last one
13101313
continue;
13111314
}
1312-
if (sig.getKey().size() < minArgCount) {
1313-
throw new IllegalArgumentException("signature " + sig.getKey() + " is missing non-optional arg for " + args);
1315+
if (sig.argTypes().size() < minArgCount) {
1316+
throw new IllegalArgumentException("signature " + sig.argTypes() + " is missing non-optional arg for " + args);
13141317
}
1315-
if (shouldHideSignature(sig.getKey(), sig.getValue())) {
1318+
if (shouldHideSignature(sig.argTypes(), sig.returnType())) {
13161319
continue;
13171320
}
13181321
builder.startObject();
13191322
builder.startArray("params");
1320-
for (int i = 0; i < sig.getKey().size(); i++) {
1323+
for (int i = 0; i < sig.argTypes().size(); i++) {
13211324
EsqlFunctionRegistry.ArgSignature arg = args.get(i);
13221325
builder.startObject();
13231326
builder.field("name", arg.name());
@@ -1332,20 +1335,20 @@ void renderKibanaFunctionDefinition(
13321335
.collect(Collectors.joining(", "))
13331336
);
13341337
} else {
1335-
builder.field("type", sig.getKey().get(i).dataType().esNameIfPossible());
1338+
builder.field("type", sig.argTypes().get(i).dataType().esNameIfPossible());
13361339
}
13371340
builder.field("optional", arg.optional());
13381341
String cleanedParamDesc = removeAppliesToBlocks(arg.description());
13391342
builder.field("description", cleanedParamDesc);
13401343
builder.endObject();
13411344
}
13421345
builder.endArray();
1343-
license = licenseChecker.invoke(sig.getKey().stream().map(Param::dataType).toList());
1346+
license = licenseChecker.invoke(sig.argTypes().stream().map(Param::dataType).toList());
13441347
if (license != null && license != License.OperationMode.BASIC) {
13451348
builder.field("license", license.toString());
13461349
}
13471350
builder.field("variadic", variadic);
1348-
builder.field("returnType", sig.getValue().esNameIfPossible());
1351+
builder.field("returnType", sig.returnType().esNameIfPossible());
13491352
builder.endObject();
13501353
}
13511354
}
@@ -1388,23 +1391,23 @@ private static String removeAppliesToBlocks(String content) {
13881391
return content.replaceAll("\\s*\\{applies_to\\}`[^`]*`\\s*", "");
13891392
}
13901393

1391-
private List<Map.Entry<List<DocsV3Support.Param>, DataType>> sortedSignatures() {
1392-
List<Map.Entry<List<DocsV3Support.Param>, DataType>> sortedSignatures = new ArrayList<>(signatures.get().entrySet());
1394+
private List<TypeSignature> sortedSignatures() {
1395+
List<TypeSignature> sortedSignatures = new ArrayList<>(signatures.get());
13931396
sortedSignatures.sort((lhs, rhs) -> {
1394-
int maxlen = Math.max(lhs.getKey().size(), rhs.getKey().size());
1397+
int maxlen = Math.max(lhs.argTypes().size(), rhs.argTypes().size());
13951398
for (int i = 0; i < maxlen; i++) {
1396-
if (lhs.getKey().size() <= i) {
1399+
if (lhs.argTypes().size() <= i) {
13971400
return -1;
13981401
}
1399-
if (rhs.getKey().size() <= i) {
1402+
if (rhs.argTypes().size() <= i) {
14001403
return 1;
14011404
}
1402-
int c = lhs.getKey().get(i).dataType().esNameIfPossible().compareTo(rhs.getKey().get(i).dataType().esNameIfPossible());
1405+
int c = lhs.argTypes().get(i).dataType().esNameIfPossible().compareTo(rhs.argTypes().get(i).dataType().esNameIfPossible());
14031406
if (c != 0) {
14041407
return c;
14051408
}
14061409
}
1407-
return lhs.getValue().esNameIfPossible().compareTo(rhs.getValue().esNameIfPossible());
1410+
return lhs.returnType().esNameIfPossible().compareTo(rhs.returnType().esNameIfPossible());
14081411
});
14091412
return sortedSignatures;
14101413
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/DocsV3SupportTests.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.HashMap;
2222
import java.util.List;
2323
import java.util.Map;
24+
import java.util.Set;
2425

2526
import static org.hamcrest.Matchers.equalTo;
2627

@@ -422,8 +423,8 @@ public TestClass(Source source, @Param(name = "str", type = { "keyword", "text"
422423
super(source, List.of(field));
423424
}
424425

425-
public static Map<List<DocsV3Support.Param>, DataType> signatures() {
426-
return Map.of(List.of(new DocsV3Support.Param(DataType.KEYWORD, List.of())), DataType.LONG);
426+
public static Set<DocsV3Support.TypeSignature> signatures() {
427+
return Set.of(new DocsV3Support.TypeSignature(List.of(new DocsV3Support.Param(DataType.KEYWORD, List.of())), DataType.LONG));
427428
}
428429

429430
@Override

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/TestClassSignatureTypes.java‎

Lines changed: 0 additions & 21 deletions
This file was deleted.

0 commit comments

Comments
 (0)