Skip to content

Commit d04b9ff

Browse files
committed
Refactor: remove some methods in tests to upper level to reduce duplication
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent bbb324d commit d04b9ff

3 files changed

Lines changed: 112 additions & 172 deletions

File tree

integ-test/src/test/java/org/opensearch/sql/legacy/OpenSearchSQLRestTestCase.java

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,18 @@
55

66
package org.opensearch.sql.legacy;
77

8+
import static org.opensearch.sql.common.setting.Settings.Key;
9+
import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
10+
import static org.opensearch.sql.legacy.TestsConstants.PERSISTENT;
11+
import static org.opensearch.sql.legacy.TestsConstants.TRANSIENT;
12+
813
import java.io.IOException;
914
import java.util.ArrayList;
1015
import java.util.List;
1116
import java.util.Map;
17+
import java.util.Objects;
1218
import java.util.Optional;
19+
import org.apache.commons.lang3.StringUtils;
1320
import org.apache.hc.client5.http.auth.AuthScope;
1421
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
1522
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
@@ -29,14 +36,19 @@
2936
import org.json.JSONArray;
3037
import org.json.JSONObject;
3138
import org.junit.AfterClass;
39+
import org.junit.Assert;
3240
import org.opensearch.client.Request;
41+
import org.opensearch.client.RequestOptions;
3342
import org.opensearch.client.Response;
3443
import org.opensearch.client.RestClient;
3544
import org.opensearch.client.RestClientBuilder;
3645
import org.opensearch.common.settings.Settings;
3746
import org.opensearch.common.unit.TimeValue;
3847
import org.opensearch.common.util.concurrent.ThreadContext;
3948
import org.opensearch.common.util.io.IOUtils;
49+
import org.opensearch.common.xcontent.XContentFactory;
50+
import org.opensearch.core.rest.RestStatus;
51+
import org.opensearch.core.xcontent.XContentBuilder;
4052
import org.opensearch.test.rest.OpenSearchRestTestCase;
4153

4254
/**
@@ -135,7 +147,13 @@ public RestClient initClient(String clusterName) throws IOException {
135147
int port = Integer.parseInt(stringUrl.substring(portSeparator + 1));
136148
hosts.add(buildHttpHost(host, port));
137149
}
138-
return buildClient(restClientSettings(), hosts.toArray(new HttpHost[0]));
150+
Settings.Builder builder = Settings.builder();
151+
if (System.getProperty("tests.rest.client_path_prefix") != null) {
152+
builder.put(CLIENT_PATH_PREFIX, System.getProperty("tests.rest.client_path_prefix"));
153+
}
154+
// Disable max compilations rate to avoid hitting compilations threshold during tests
155+
builder.put(Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE.getKeyValue(), "true");
156+
return buildClient(builder.build(), hosts.toArray(new HttpHost[0]));
139157
}
140158

141159
/** Get a comma delimited list of [host:port] to which to send REST requests. */
@@ -301,4 +319,97 @@ public void configureMultiClusters(String remote) throws IOException {
301319
connectionRequest.setJsonEntity(connectionSetting);
302320
adminClient().performRequest(connectionRequest);
303321
}
322+
323+
/**
324+
* Increases the maximum script compilation rate for all script contexts for tests. This method
325+
* sets an unlimited compilation rate for each script context when the
326+
* script.disable_max_compilations_rate setting is not enabled, allowing tests to run without
327+
* hitting compilation rate limits.
328+
*
329+
* @throws IOException if there is an error retrieving cluster settings or updating them
330+
*/
331+
protected void increaseMaxCompilationsRate() throws IOException {
332+
// When script.disable_max_compilations_rate is set, custom context compilation rates cannot be
333+
// set
334+
if (!Objects.equals(
335+
getClusterSetting(
336+
org.opensearch.sql.common.setting.Settings.Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE
337+
.getKeyValue(),
338+
"persistent"),
339+
"true")) {
340+
List<String> contexts = getScriptContexts();
341+
for (String context : contexts) {
342+
String contextCompilationsRate =
343+
org.opensearch.sql.common.setting.Settings.Key
344+
.SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN
345+
.getKeyValue()
346+
.replace("*", context);
347+
updateClusterSetting(contextCompilationsRate, "unlimited", true);
348+
}
349+
}
350+
}
351+
352+
protected List<String> getScriptContexts() throws IOException {
353+
Request request = new Request("GET", "/_script_context");
354+
String responseBody = executeRequest(request);
355+
JSONObject jsonResponse = new JSONObject(responseBody);
356+
JSONArray contexts = jsonResponse.getJSONArray("contexts");
357+
List<String> contextNames = new ArrayList<>();
358+
for (int i = 0; i < contexts.length(); i++) {
359+
JSONObject context = contexts.getJSONObject(i);
360+
String contextName = context.getString("name");
361+
contextNames.add(contextName);
362+
}
363+
return contextNames;
364+
}
365+
366+
protected static void updateClusterSetting(String settingKey, Object value) throws IOException {
367+
updateClusterSetting(settingKey, value, true);
368+
}
369+
370+
protected static void updateClusterSetting(String settingKey, Object value, boolean persistent)
371+
throws IOException {
372+
String property = persistent ? PERSISTENT : TRANSIENT;
373+
XContentBuilder builder =
374+
XContentFactory.jsonBuilder()
375+
.startObject()
376+
.startObject(property)
377+
.field(settingKey, value)
378+
.endObject()
379+
.endObject();
380+
Request request = new Request("PUT", "_cluster/settings");
381+
request.setJsonEntity(builder.toString());
382+
Response response = client().performRequest(request);
383+
Assert.assertEquals(
384+
RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
385+
}
386+
387+
protected static JSONObject getAllClusterSettings() throws IOException {
388+
Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults");
389+
RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
390+
restOptionsBuilder.addHeader("Content-Type", "application/json");
391+
request.setOptions(restOptionsBuilder);
392+
return new JSONObject(executeRequest(request));
393+
}
394+
395+
protected static String getClusterSetting(String settingPath, String type) throws IOException {
396+
JSONObject settings = getAllClusterSettings();
397+
String value = settings.optJSONObject(type).optString(settingPath);
398+
if (StringUtils.isEmpty(value)) {
399+
return settings.optJSONObject("defaults").optString(settingPath);
400+
} else {
401+
return value;
402+
}
403+
}
404+
405+
protected static String executeRequest(final Request request) throws IOException {
406+
return executeRequest(request, client());
407+
}
408+
409+
protected static String executeRequest(final Request request, RestClient client)
410+
throws IOException {
411+
Response response = client.performRequest(request);
412+
Assert.assertEquals(200, response.getStatusLine().getStatusCode());
413+
return getResponseBody(response);
414+
}
304415
}

integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java

Lines changed: 0 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -23,38 +23,21 @@
2323
import static org.opensearch.sql.legacy.TestUtils.getOrderIndexMapping;
2424
import static org.opensearch.sql.legacy.TestUtils.getPeople2IndexMapping;
2525
import static org.opensearch.sql.legacy.TestUtils.getPhraseIndexMapping;
26-
import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
2726
import static org.opensearch.sql.legacy.TestUtils.getWeblogsIndexMapping;
2827
import static org.opensearch.sql.legacy.TestUtils.isIndexExist;
2928
import static org.opensearch.sql.legacy.TestUtils.loadDataByRestClient;
30-
import static org.opensearch.sql.legacy.TestsConstants.PERSISTENT;
31-
import static org.opensearch.sql.legacy.TestsConstants.TRANSIENT;
3229

3330
import java.io.IOException;
3431
import java.nio.file.Files;
3532
import java.nio.file.Path;
3633
import java.nio.file.Paths;
37-
import java.util.ArrayList;
38-
import java.util.List;
39-
import java.util.Objects;
4034
import javax.management.MBeanServerInvocationHandler;
4135
import javax.management.ObjectName;
4236
import javax.management.remote.JMXConnector;
4337
import javax.management.remote.JMXConnectorFactory;
4438
import javax.management.remote.JMXServiceURL;
45-
import org.apache.commons.lang3.StringUtils;
46-
import org.json.JSONArray;
47-
import org.json.JSONObject;
4839
import org.junit.AfterClass;
49-
import org.junit.Assert;
5040
import org.junit.Before;
51-
import org.opensearch.client.Request;
52-
import org.opensearch.client.RequestOptions;
53-
import org.opensearch.client.Response;
54-
import org.opensearch.common.xcontent.XContentFactory;
55-
import org.opensearch.core.rest.RestStatus;
56-
import org.opensearch.core.xcontent.XContentBuilder;
57-
import org.opensearch.sql.common.setting.Settings;
5841

5942
/**
6043
*
@@ -161,87 +144,6 @@ protected void init() throws Exception {
161144
increaseMaxCompilationsRate();
162145
}
163146

164-
protected static void updateClusterSetting(String settingKey, Object value) throws IOException {
165-
updateClusterSetting(settingKey, value, true);
166-
}
167-
168-
protected static void updateClusterSetting(String settingKey, Object value, boolean persistent)
169-
throws IOException {
170-
String property = persistent ? PERSISTENT : TRANSIENT;
171-
XContentBuilder builder =
172-
XContentFactory.jsonBuilder()
173-
.startObject()
174-
.startObject(property)
175-
.field(settingKey, value)
176-
.endObject()
177-
.endObject();
178-
Request request = new Request("PUT", "_cluster/settings");
179-
request.setJsonEntity(builder.toString());
180-
Response response = client().performRequest(request);
181-
Assert.assertEquals(
182-
RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
183-
}
184-
185-
/**
186-
* Increases the maximum script compilation rate for all script contexts for tests. This method
187-
* sets an unlimited compilation rate for each script context when the
188-
* script.disable_max_compilations_rate setting is not enabled, allowing tests to run without
189-
* hitting compilation rate limits.
190-
*
191-
* @throws IOException if there is an error retrieving cluster settings or updating them
192-
*/
193-
protected void increaseMaxCompilationsRate() throws IOException {
194-
// When script.disable_max_compilations_rate is set, custom context compilation rates cannot be
195-
// set
196-
if (!Objects.equals(
197-
getClusterSetting(
198-
Settings.Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE.getKeyValue(), "persistent"),
199-
"true")) {
200-
List<String> contexts = getScriptContexts();
201-
for (String context : contexts) {
202-
String contextCompilationsRate =
203-
Settings.Key.SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN
204-
.getKeyValue()
205-
.replace("*", context);
206-
updateClusterSetting(contextCompilationsRate, "unlimited", true);
207-
}
208-
}
209-
}
210-
211-
protected List<String> getScriptContexts() throws IOException {
212-
Request request = new Request("GET", "/_script_context");
213-
Response response = client().performRequest(request);
214-
String responseBody = getResponseBody(response);
215-
JSONObject jsonResponse = new JSONObject(responseBody);
216-
JSONArray contexts = jsonResponse.getJSONArray("contexts");
217-
List<String> contextNames = new ArrayList<>();
218-
for (int i = 0; i < contexts.length(); i++) {
219-
JSONObject context = contexts.getJSONObject(i);
220-
String contextName = context.getString("name");
221-
contextNames.add(contextName);
222-
}
223-
return contextNames;
224-
}
225-
226-
protected static String getClusterSetting(String settingPath, String type) throws IOException {
227-
JSONObject settings = getAllClusterSettings();
228-
String value = settings.optJSONObject(type).optString(settingPath);
229-
if (StringUtils.isEmpty(value)) {
230-
return settings.optJSONObject("defaults").optString(settingPath);
231-
} else {
232-
return value;
233-
}
234-
}
235-
236-
protected static JSONObject getAllClusterSettings() throws IOException {
237-
Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults");
238-
RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
239-
restOptionsBuilder.addHeader("Content-Type", "application/json");
240-
request.setOptions(restOptionsBuilder);
241-
Response response = client().performRequest(request);
242-
return new JSONObject(getResponseBody(response));
243-
}
244-
245147
protected static void wipeAllClusterSettings() throws IOException {
246148
updateClusterSetting("*", null, true);
247149
updateClusterSetting("*", null, false);

integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java

Lines changed: 0 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
package org.opensearch.sql.legacy;
77

88
import static com.google.common.base.Strings.isNullOrEmpty;
9-
import static org.opensearch.sql.common.setting.Settings.Key.SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN;
109
import static org.opensearch.sql.legacy.TestUtils.createIndexByRestClient;
1110
import static org.opensearch.sql.legacy.TestUtils.getAccountIndexMapping;
1211
import static org.opensearch.sql.legacy.TestUtils.getAliasIndexMapping;
@@ -64,11 +63,8 @@
6463
import java.nio.file.Files;
6564
import java.nio.file.Path;
6665
import java.nio.file.Paths;
67-
import java.util.ArrayList;
68-
import java.util.List;
6966
import java.util.Locale;
7067
import java.util.Map;
71-
import java.util.Objects;
7268
import java.util.concurrent.Callable;
7369
import javax.management.MBeanServerInvocationHandler;
7470
import javax.management.ObjectName;
@@ -222,46 +218,6 @@ protected void resetMaxResultWindow(String indexName) throws IOException {
222218
indexName, "{ \"index\": { \"max_result_window\": " + DEFAULT_MAX_RESULT_WINDOW + " } }");
223219
}
224220

225-
/**
226-
* Increases the maximum script compilation rate for all script contexts for tests. This method
227-
* sets an unlimited compilation rate for each script context when the
228-
* script.disable_max_compilations_rate setting is not enabled, allowing tests to run without
229-
* hitting compilation rate limits.
230-
*
231-
* @throws IOException if there is an error retrieving cluster settings or updating them
232-
*/
233-
protected void increaseMaxCompilationsRate() throws IOException {
234-
// When script.disable_max_compilations_rate is set, custom context compilation rates cannot be
235-
// set
236-
if (!Objects.equals(
237-
getClusterSetting(
238-
Settings.Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE.getKeyValue(), "persistent"),
239-
"true")) {
240-
List<String> contexts = getScriptContexts();
241-
for (String context : contexts) {
242-
String contextCompilationsRate =
243-
SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN.getKeyValue().replace("*", context);
244-
updateClusterSettings(
245-
new ClusterSetting("persistent", contextCompilationsRate, "unlimited"));
246-
}
247-
}
248-
}
249-
250-
protected List<String> getScriptContexts() throws IOException {
251-
Request request = new Request("GET", "/_script_context");
252-
Response response = client().performRequest(request);
253-
String responseBody = getResponseBody(response);
254-
JSONObject jsonResponse = new JSONObject(responseBody);
255-
JSONArray contexts = jsonResponse.getJSONArray("contexts");
256-
List<String> contextNames = new ArrayList<>();
257-
for (int i = 0; i < contexts.length(); i++) {
258-
JSONObject context = contexts.getJSONObject(i);
259-
String contextName = context.getString("name");
260-
contextNames.add(contextName);
261-
}
262-
return contextNames;
263-
}
264-
265221
/** Provide for each test to load test index, data and other setup work */
266222
protected void init() throws Exception {
267223
disableCalcite();
@@ -442,17 +398,6 @@ private String executeRequest(final String requestBody, final boolean isExplainQ
442398
return executeRequest(sqlRequest);
443399
}
444400

445-
protected static String executeRequest(final Request request, RestClient client)
446-
throws IOException {
447-
Response response = client.performRequest(request);
448-
Assert.assertEquals(200, response.getStatusLine().getStatusCode());
449-
return getResponseBody(response);
450-
}
451-
452-
protected static String executeRequest(final Request request) throws IOException {
453-
return executeRequest(request, client());
454-
}
455-
456401
protected JSONObject executeQueryWithGetRequest(final String sqlQuery) throws IOException {
457402

458403
final Request request = buildGetEndpointRequest(sqlQuery);
@@ -489,24 +434,6 @@ protected static JSONObject updateClusterSettings(ClusterSetting setting) throws
489434
return updateClusterSettings(setting, client());
490435
}
491436

492-
protected static JSONObject getAllClusterSettings() throws IOException {
493-
Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults");
494-
RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
495-
restOptionsBuilder.addHeader("Content-Type", "application/json");
496-
request.setOptions(restOptionsBuilder);
497-
return new JSONObject(executeRequest(request));
498-
}
499-
500-
protected static String getClusterSetting(String settingPath, String type) throws IOException {
501-
JSONObject settings = getAllClusterSettings();
502-
String value = settings.optJSONObject(type).optString(settingPath);
503-
if (StringUtils.isEmpty(value)) {
504-
return settings.optJSONObject("defaults").optString(settingPath);
505-
} else {
506-
return value;
507-
}
508-
}
509-
510437
protected static class ClusterSetting {
511438
private final String type;
512439
private final String name;

0 commit comments

Comments
 (0)