Skip to content

Commit 1432c79

Browse files
alex-spiesalex-spies
committed
Csv loading in multi-cluster test via ClassRule
1 parent 52f6789 commit 1432c79

File tree

5 files changed

+99
-56
lines changed

5 files changed

+99
-56
lines changed

‎x-pack/plugin/esql/qa/server/mixed-cluster/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/mixed/MixedClusterEsqlSpecIT.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,28 @@
88
package org.elasticsearch.xpack.esql.qa.mixed;
99

1010
import org.elasticsearch.Version;
11+
import org.elasticsearch.client.RestClient;
12+
import org.elasticsearch.common.settings.Settings;
1113
import org.elasticsearch.test.cluster.ElasticsearchCluster;
1214
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
1315
import org.elasticsearch.xpack.ql.CsvSpecReader.CsvTestCase;
1416
import org.junit.ClassRule;
1517
import org.junit.rules.RuleChain;
1618
import org.junit.rules.TestRule;
1719

20+
import java.io.IOException;
21+
1822
import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
1923
import static org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase.Mode.ASYNC;
2024

2125
public class MixedClusterEsqlSpecIT extends EsqlSpecTestCase {
2226
public static ElasticsearchCluster cluster = Clusters.mixedVersionCluster();
23-
public static TestRuleRestClient client = new TestRuleRestClient(cluster);
27+
public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
28+
@Override
29+
protected RestClient provideObject() throws IOException {
30+
return startClient(cluster, Settings.builder().build());
31+
}
32+
};
2433
public static CsvLoader loader = new CsvLoader(client);
2534

2635
@ClassRule

‎x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
2222
import org.elasticsearch.xpack.ql.CsvSpecReader;
2323
import org.elasticsearch.xpack.ql.CsvSpecReader.CsvTestCase;
24-
import org.junit.Before;
2524
import org.junit.ClassRule;
2625
import org.junit.rules.RuleChain;
2726
import org.junit.rules.TestRule;
@@ -35,9 +34,7 @@
3534
import java.util.stream.Collectors;
3635

3736
import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
38-
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.CSV_DATASET_MAP;
3937
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.ENRICH_SOURCE_INDICES;
40-
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
4138
import static org.mockito.ArgumentMatchers.any;
4239
import static org.mockito.Mockito.doAnswer;
4340
import static org.mockito.Mockito.mock;
@@ -53,16 +50,17 @@ public class MultiClusterSpecIT extends EsqlSpecTestCase {
5350

5451
static ElasticsearchCluster remoteCluster = Clusters.remoteCluster();
5552
static ElasticsearchCluster localCluster = Clusters.localCluster(remoteCluster);
53+
public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
54+
@Override
55+
protected RestClient provideObject() throws IOException {
56+
HttpHost[] localHosts = parseClusterHostsStatic(localCluster.getHttpAddresses()).toArray(HttpHost[]::new);
57+
return doBuildClient(Settings.builder().build(), localHosts);
58+
}
59+
};
60+
public static CsvLoader loader = new CsvLoader(client);
5661

5762
@ClassRule
58-
public static TestRule clusterRule = RuleChain.outerRule(remoteCluster).around(localCluster);
59-
60-
@Before
61-
public void setup() throws IOException {
62-
if (indexExists(CSV_DATASET_MAP.keySet().iterator().next()) == false) {
63-
loadDataSetIntoEs(client());
64-
}
65-
}
63+
public static TestRule clusterRule = RuleChain.outerRule(remoteCluster).around(localCluster).around(client).around(loader);
6664

6765
public MultiClusterSpecIT(String fileName, String groupName, String testName, Integer lineNumber, CsvTestCase testCase, Mode mode) {
6866
super(fileName, groupName, testName, lineNumber, convertToRemoteIndices(testCase), mode);
@@ -82,9 +80,13 @@ protected String getTestRestCluster() {
8280

8381
@Override
8482
protected RestClient buildClient(Settings settings, HttpHost[] localHosts) throws IOException {
85-
RestClient localClient = super.buildClient(settings, localHosts);
86-
HttpHost[] remoteHosts = parseClusterHosts(remoteCluster.getHttpAddresses()).toArray(HttpHost[]::new);
87-
RestClient remoteClient = super.buildClient(settings, remoteHosts);
83+
return doBuildClient(settings, localHosts);
84+
}
85+
86+
private static RestClient doBuildClient(Settings settings, HttpHost[] localHosts) throws IOException {
87+
RestClient localClient = buildClientStatic(settings, localHosts);
88+
HttpHost[] remoteHosts = parseClusterHostsStatic(remoteCluster.getHttpAddresses()).toArray(HttpHost[]::new);
89+
RestClient remoteClient = buildClientStatic(settings, remoteHosts);
8890
return twoClients(localClient, remoteClient);
8991
}
9092

‎x-pack/plugin/esql/qa/server/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/multi_node/EsqlSpecIT.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,25 @@
77

88
package org.elasticsearch.xpack.esql.qa.multi_node;
99

10+
import org.elasticsearch.client.RestClient;
11+
import org.elasticsearch.common.settings.Settings;
1012
import org.elasticsearch.test.cluster.ElasticsearchCluster;
1113
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
1214
import org.elasticsearch.xpack.ql.CsvSpecReader.CsvTestCase;
1315
import org.junit.ClassRule;
1416
import org.junit.rules.RuleChain;
1517
import org.junit.rules.TestRule;
1618

19+
import java.io.IOException;
20+
1721
public class EsqlSpecIT extends EsqlSpecTestCase {
1822
public static ElasticsearchCluster cluster = Clusters.testCluster();
19-
public static TestRuleRestClient client = new TestRuleRestClient(cluster);
23+
public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
24+
@Override
25+
protected RestClient provideObject() throws IOException {
26+
return startClient(cluster, Settings.builder().build());
27+
}
28+
};
2029
public static CsvLoader loader = new CsvLoader(client);
2130

2231
@ClassRule

‎x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/EsqlSpecIT.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99

1010
import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters;
1111

12+
import org.elasticsearch.client.RestClient;
13+
import org.elasticsearch.common.settings.Settings;
1214
import org.elasticsearch.test.TestClustersThreadFilter;
1315
import org.elasticsearch.test.cluster.ElasticsearchCluster;
1416
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
@@ -17,10 +19,17 @@
1719
import org.junit.rules.RuleChain;
1820
import org.junit.rules.TestRule;
1921

22+
import java.io.IOException;
23+
2024
@ThreadLeakFilters(filters = TestClustersThreadFilter.class)
2125
public class EsqlSpecIT extends EsqlSpecTestCase {
2226
public static ElasticsearchCluster cluster = Clusters.testCluster();
23-
public static TestRuleRestClient client = new TestRuleRestClient(cluster);
27+
public static ClosingTestRule<RestClient> client = new ClosingTestRule<>() {
28+
@Override
29+
protected RestClient provideObject() throws IOException {
30+
return startClient(cluster, Settings.builder().build());
31+
}
32+
};
2433
public static CsvLoader loader = new CsvLoader(client);
2534

2635
@ClassRule

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/EsqlSpecTestCase.java‎

Lines changed: 53 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -38,13 +38,15 @@
3838
import org.junit.runner.Description;
3939
import org.junit.runners.model.Statement;
4040

41+
import java.io.Closeable;
4142
import java.io.IOException;
4243
import java.net.URL;
4344
import java.util.ArrayList;
4445
import java.util.List;
4546
import java.util.Locale;
4647
import java.util.Map;
4748

49+
import static java.util.Collections.unmodifiableList;
4850
import static org.apache.lucene.geo.GeoEncodingUtils.decodeLatitude;
4951
import static org.apache.lucene.geo.GeoEncodingUtils.decodeLongitude;
5052
import static org.apache.lucene.geo.GeoEncodingUtils.encodeLatitude;
@@ -242,10 +244,51 @@ public static void assertRequestBreakerEmpty() throws Exception {
242244
});
243245
}
244246

247+
public static RestClient startClient(ElasticsearchCluster cluster, Settings clientSettings) throws IOException {
248+
String address = cluster.getHttpAddress(0);
249+
int portSeparator = address.lastIndexOf(':');
250+
if (portSeparator < 0) {
251+
throw new IllegalArgumentException("Illegal cluster url [" + address + "]");
252+
}
253+
String host = address.substring(0, portSeparator);
254+
int port = Integer.parseInt(address.substring(portSeparator + 1));
255+
HttpHost[] httpHosts = new HttpHost[] { new HttpHost(host, port) };
256+
257+
return buildClientStatic(clientSettings, httpHosts);
258+
}
259+
260+
/**
261+
* Like {@link ESRestTestCase#buildClient(Settings, HttpHost[])} but static.
262+
*/
263+
public static RestClient buildClientStatic(Settings settings, HttpHost[] hosts) throws IOException {
264+
RestClientBuilder builder = RestClient.builder(hosts);
265+
doConfigureClient(builder, settings);
266+
builder.setStrictDeprecationMode(true);
267+
return builder.build();
268+
}
269+
270+
/**
271+
* Like {@link ESRestTestCase#parseClusterHosts(String)} but static.
272+
*/
273+
protected static List<HttpHost> parseClusterHostsStatic(String hostsString) {
274+
String[] stringUrls = hostsString.split(",");
275+
List<HttpHost> hosts = new ArrayList<>(stringUrls.length);
276+
for (String stringUrl : stringUrls) {
277+
int portSeparator = stringUrl.lastIndexOf(':');
278+
if (portSeparator < 0) {
279+
throw new IllegalArgumentException("Illegal cluster url [" + stringUrl + "]");
280+
}
281+
String host = stringUrl.substring(0, portSeparator);
282+
int port = Integer.valueOf(stringUrl.substring(portSeparator + 1));
283+
hosts.add(new HttpHost(host, port, "http"));
284+
}
285+
return unmodifiableList(hosts);
286+
}
287+
245288
public static class CsvLoader implements TestRule {
246-
TestRuleRestClient client;
289+
ClosingTestRule<RestClient> client;
247290

248-
public CsvLoader(TestRuleRestClient client) {
291+
public CsvLoader(ClosingTestRule<RestClient> client) {
249292
this.client = client;
250293
}
251294

@@ -254,64 +297,35 @@ public Statement apply(Statement base, Description description) {
254297
return new Statement() {
255298
@Override
256299
public void evaluate() throws Throwable {
257-
loadDataSetIntoEs(client.client());
300+
loadDataSetIntoEs(client.get());
258301
base.evaluate();
259302
}
260303
};
261304
}
262305
}
263306

264-
public static class TestRuleRestClient implements TestRule {
265-
private ElasticsearchCluster cluster;
266-
private Settings clientSettings;
267-
private RestClient client;
307+
public abstract static class ClosingTestRule<T extends Closeable> implements TestRule {
308+
private T providedObject;
268309

269-
public TestRuleRestClient(ElasticsearchCluster cluster, Settings clientSettings) {
270-
this.cluster = cluster;
271-
this.clientSettings = clientSettings;
310+
public T get() {
311+
return providedObject;
272312
}
273313

274-
public TestRuleRestClient(ElasticsearchCluster cluster) {
275-
this(cluster, Settings.builder().build());
276-
}
277-
278-
public RestClient client() {
279-
return client;
280-
}
314+
protected abstract T provideObject() throws IOException;
281315

282316
@Override
283317
public Statement apply(Statement base, Description description) {
284318
return new Statement() {
285319
@Override
286320
public void evaluate() throws Throwable {
287321
try {
288-
client = doStartClient();
322+
providedObject = provideObject();
289323
base.evaluate();
290324
} finally {
291-
IOUtils.close(client);
325+
IOUtils.close(providedObject);
292326
}
293327
}
294328
};
295329
}
296-
297-
private RestClient doStartClient() throws IOException {
298-
String address = cluster.getHttpAddress(0);
299-
int portSeparator = address.lastIndexOf(':');
300-
if (portSeparator < 0) {
301-
throw new IllegalArgumentException("Illegal cluster url [" + address + "]");
302-
}
303-
String host = address.substring(0, portSeparator);
304-
int port = Integer.parseInt(address.substring(portSeparator + 1));
305-
HttpHost[] httpHosts = new HttpHost[] { new HttpHost(host, port) };
306-
307-
return buildClient(clientSettings, httpHosts);
308-
}
309-
310-
private static RestClient buildClient(Settings settings, HttpHost[] hosts) throws IOException {
311-
RestClientBuilder builder = RestClient.builder(hosts);
312-
doConfigureClient(builder, settings);
313-
builder.setStrictDeprecationMode(true);
314-
return builder.build();
315-
}
316330
}
317331
}

0 commit comments

Comments
 (0)