Skip to content

Commit 60c4c6c

Browse files
RyanL1997RyanL1997
committed
Add tests
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
1 parent a4d156e commit 60c4c6c

7 files changed

Lines changed: 529 additions & 0 deletions

File tree

‎core/src/test/java/org/opensearch/sql/executor/execution/ExplainPlanTest.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,22 @@
66
package org.opensearch.sql.executor.execution;
77

88
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertNull;
910
import static org.junit.jupiter.api.Assertions.assertThrows;
1011
import static org.mockito.ArgumentMatchers.any;
12+
import static org.mockito.Mockito.doAnswer;
1113
import static org.mockito.Mockito.doNothing;
1214
import static org.mockito.Mockito.times;
1315
import static org.mockito.Mockito.verify;
1416

17+
import java.util.Map;
18+
import java.util.concurrent.atomic.AtomicReference;
1519
import org.junit.jupiter.api.Test;
1620
import org.junit.jupiter.api.extension.ExtendWith;
1721
import org.mockito.Mock;
1822
import org.mockito.junit.jupiter.MockitoExtension;
1923
import org.opensearch.sql.ast.statement.ExplainMode;
24+
import org.opensearch.sql.calcite.CalcitePlanContext;
2025
import org.opensearch.sql.common.response.ResponseListener;
2126
import org.opensearch.sql.executor.ExecutionEngine;
2227
import org.opensearch.sql.executor.QueryId;
@@ -56,4 +61,36 @@ public void explainThrowException() {
5661
});
5762
assertEquals("explain query can not been explained.", unsupportedExplainException.getMessage());
5863
}
64+
65+
@Test
66+
public void execute_sets_highlight_threadlocal_for_explain() {
67+
Map<String, Object> highlightConfig = Map.of("fields", Map.of("*", Map.of()));
68+
AtomicReference<Map<String, Object>> captured = new AtomicReference<>();
69+
70+
doAnswer(
71+
invocation -> {
72+
captured.set(CalcitePlanContext.getHighlightConfig());
73+
return null;
74+
})
75+
.when(queryPlan)
76+
.explain(any(), any());
77+
78+
ExplainPlan explainPlan = new ExplainPlan(queryId, queryType, queryPlan, mode, explainListener);
79+
explainPlan.setHighlightConfig(highlightConfig);
80+
explainPlan.execute();
81+
82+
assertEquals(highlightConfig, captured.get());
83+
}
84+
85+
@Test
86+
public void execute_clears_highlight_threadlocal_after_explain() {
87+
Map<String, Object> highlightConfig = Map.of("fields", Map.of("*", Map.of()));
88+
doNothing().when(queryPlan).explain(any(), any());
89+
90+
ExplainPlan explainPlan = new ExplainPlan(queryId, queryType, queryPlan, mode, explainListener);
91+
explainPlan.setHighlightConfig(highlightConfig);
92+
explainPlan.execute();
93+
94+
assertNull(CalcitePlanContext.getHighlightConfig());
95+
}
5996
}

‎core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,19 @@
55

66
package org.opensearch.sql.executor.execution;
77

8+
import static org.junit.jupiter.api.Assertions.assertEquals;
89
import static org.junit.jupiter.api.Assertions.assertNotNull;
10+
import static org.junit.jupiter.api.Assertions.assertNull;
911
import static org.junit.jupiter.api.Assertions.assertTrue;
1012
import static org.junit.jupiter.api.Assertions.fail;
1113
import static org.mockito.ArgumentMatchers.any;
14+
import static org.mockito.Mockito.doAnswer;
1215
import static org.mockito.Mockito.mock;
1316
import static org.mockito.Mockito.times;
1417
import static org.mockito.Mockito.verify;
1518

19+
import java.util.Map;
20+
import java.util.concurrent.atomic.AtomicReference;
1621
import org.apache.commons.lang3.NotImplementedException;
1722
import org.junit.jupiter.api.DisplayNameGeneration;
1823
import org.junit.jupiter.api.DisplayNameGenerator;
@@ -22,6 +27,7 @@
2227
import org.mockito.junit.jupiter.MockitoExtension;
2328
import org.opensearch.sql.ast.statement.ExplainMode;
2429
import org.opensearch.sql.ast.tree.UnresolvedPlan;
30+
import org.opensearch.sql.calcite.CalcitePlanContext;
2531
import org.opensearch.sql.common.response.ResponseListener;
2632
import org.opensearch.sql.executor.DefaultExecutionEngine;
2733
import org.opensearch.sql.executor.ExecutionEngine;
@@ -126,4 +132,35 @@ public void onFailure(Exception e) {
126132
},
127133
mode);
128134
}
135+
136+
@Test
137+
public void execute_sets_highlight_threadlocal() {
138+
Map<String, Object> highlightConfig = Map.of("fields", Map.of("*", Map.of()));
139+
AtomicReference<Map<String, Object>> captured = new AtomicReference<>();
140+
141+
doAnswer(
142+
invocation -> {
143+
captured.set(CalcitePlanContext.getHighlightConfig());
144+
return null;
145+
})
146+
.when(queryService)
147+
.execute(any(), any(), any());
148+
149+
QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
150+
query.setHighlightConfig(highlightConfig);
151+
query.execute();
152+
153+
assertEquals(highlightConfig, captured.get());
154+
}
155+
156+
@Test
157+
public void execute_clears_highlight_threadlocal_after_execution() {
158+
Map<String, Object> highlightConfig = Map.of("fields", Map.of("*", Map.of()));
159+
160+
QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
161+
query.setHighlightConfig(highlightConfig);
162+
query.execute();
163+
164+
assertNull(CalcitePlanContext.getHighlightConfig());
165+
}
129166
}

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLHighlightIT.java‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.remote;
7+
8+
import static org.junit.Assert.assertEquals;
9+
import static org.junit.Assert.assertFalse;
10+
import static org.junit.Assert.assertNotNull;
11+
import static org.junit.Assert.assertTrue;
12+
13+
import java.io.IOException;
14+
import java.util.Locale;
15+
import org.json.JSONArray;
16+
import org.json.JSONObject;
17+
import org.junit.jupiter.api.Test;
18+
import org.opensearch.client.Request;
19+
import org.opensearch.client.RequestOptions;
20+
import org.opensearch.client.Response;
21+
import org.opensearch.sql.ppl.PPLIntegTestCase;
22+
23+
public class CalcitePPLHighlightIT extends PPLIntegTestCase {
24+
25+
private static final String TEST_INDEX = "highlight_test";
26+
27+
@Override
28+
public void init() throws Exception {
29+
super.init();
30+
enableCalcite();
31+
32+
// Create index with text fields
33+
Request createIndex = new Request("PUT", "/" + TEST_INDEX);
34+
createIndex.setJsonEntity(
35+
"{"
36+
+ "\"settings\": {\"number_of_shards\": 1, \"number_of_replicas\": 0},"
37+
+ "\"mappings\": {\"properties\": {"
38+
+ "\"message\": {\"type\": \"text\"},"
39+
+ "\"status\": {\"type\": \"text\"},"
40+
+ "\"code\": {\"type\": \"integer\"}"
41+
+ "}}"
42+
+ "}");
43+
client().performRequest(createIndex);
44+
45+
// Index test documents
46+
Request bulk = new Request("POST", "/" + TEST_INDEX + "/_bulk?refresh=true");
47+
bulk.setJsonEntity(
48+
"{\"index\": {}}\n"
49+
+ "{\"message\": \"Connection error occurred\", \"status\": \"error response\","
50+
+ " \"code\": 500}\n"
51+
+ "{\"index\": {}}\n"
52+
+ "{\"message\": \"Request completed successfully\", \"status\": \"ok\", \"code\":"
53+
+ " 200}\n"
54+
+ "{\"index\": {}}\n"
55+
+ "{\"message\": \"Timeout error in service\", \"status\": \"error timeout\", \"code\":"
56+
+ " 504}\n");
57+
client().performRequest(bulk);
58+
}
59+
60+
@Test
61+
public void testHighlightWithWildcardFields() throws IOException {
62+
JSONObject result =
63+
executeQueryWithHighlight(
64+
"search source=" + TEST_INDEX + " \"error\"",
65+
"{\"fields\": {\"*\": {}}, \"pre_tags\": [\"<em>\"], \"post_tags\": [\"</em>\"]}");
66+
67+
assertTrue(result.has("highlights"));
68+
JSONArray highlights = result.getJSONArray("highlights");
69+
assertEquals(result.getInt("size"), highlights.length());
70+
71+
// At least one highlight entry should contain "error" wrapped in tags
72+
boolean foundHighlight = false;
73+
for (int i = 0; i < highlights.length(); i++) {
74+
if (!highlights.isNull(i)) {
75+
String hlStr = highlights.get(i).toString();
76+
if (hlStr.contains("<em>error</em>") || hlStr.contains("<em>Error</em>")) {
77+
foundHighlight = true;
78+
break;
79+
}
80+
}
81+
}
82+
assertTrue("Expected at least one highlight with <em> tags", foundHighlight);
83+
}
84+
85+
@Test
86+
public void testHighlightWithSpecificField() throws IOException {
87+
JSONObject result =
88+
executeQueryWithHighlight(
89+
"search source=" + TEST_INDEX + " \"error\"",
90+
"{\"fields\": {\"message\": {}}, \"pre_tags\": [\"<em>\"], \"post_tags\":"
91+
+ " [\"</em>\"]}");
92+
93+
assertTrue(result.has("highlights"));
94+
JSONArray highlights = result.getJSONArray("highlights");
95+
96+
// Check that highlights only contain "message" field, not "status"
97+
for (int i = 0; i < highlights.length(); i++) {
98+
if (!highlights.isNull(i)) {
99+
JSONObject hl = highlights.getJSONObject(i);
100+
assertFalse("Should not highlight status field", hl.has("status"));
101+
}
102+
}
103+
}
104+
105+
@Test
106+
public void testHighlightWithCustomTags() throws IOException {
107+
JSONObject result =
108+
executeQueryWithHighlight(
109+
"search source=" + TEST_INDEX + " \"error\"",
110+
"{\"fields\": {\"*\": {}}, \"pre_tags\": [\"<mark>\"], \"post_tags\": [\"</mark>\"]}");
111+
112+
assertTrue(result.has("highlights"));
113+
JSONArray highlights = result.getJSONArray("highlights");
114+
115+
boolean foundCustomTag = false;
116+
for (int i = 0; i < highlights.length(); i++) {
117+
if (!highlights.isNull(i)) {
118+
String hlStr = highlights.get(i).toString();
119+
if (hlStr.contains("<mark>")) {
120+
foundCustomTag = true;
121+
break;
122+
}
123+
}
124+
}
125+
assertTrue("Expected custom <mark> tags in highlights", foundCustomTag);
126+
}
127+
128+
@Test
129+
public void testNoHighlightWhenNotRequested() throws IOException {
130+
JSONObject result = executeQuery("search source=" + TEST_INDEX + " \"error\"");
131+
132+
assertFalse("Should not have highlights when not requested", result.has("highlights"));
133+
}
134+
135+
@Test
136+
public void testHighlightWithPipedFilter() throws IOException {
137+
JSONObject result =
138+
executeQueryWithHighlight(
139+
"search source=" + TEST_INDEX + " \"error\" | where code > 500",
140+
"{\"fields\": {\"*\": {}}, \"pre_tags\": [\"<em>\"], \"post_tags\": [\"</em>\"]}");
141+
142+
assertTrue(result.has("highlights"));
143+
// Only the doc with code=504 should match
144+
assertEquals(1, result.getInt("size"));
145+
JSONArray highlights = result.getJSONArray("highlights");
146+
assertEquals(1, highlights.length());
147+
assertNotNull(highlights.get(0));
148+
}
149+
150+
@Test
151+
public void testExplainWithHighlight() throws IOException {
152+
Request request = new Request("POST", "/_plugins/_ppl/_explain");
153+
request.setJsonEntity(
154+
String.format(
155+
Locale.ROOT,
156+
"{\"query\": \"search source=%s \\\"error\\\"\","
157+
+ "\"highlight\": {\"fields\": {\"*\": {}},"
158+
+ "\"pre_tags\": [\"<em>\"], \"post_tags\": [\"</em>\"]}}",
159+
TEST_INDEX));
160+
RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
161+
restOptionsBuilder.addHeader("Content-Type", "application/json");
162+
request.setOptions(restOptionsBuilder);
163+
Response response = client().performRequest(request);
164+
assertEquals(200, response.getStatusLine().getStatusCode());
165+
166+
String body = org.opensearch.sql.legacy.TestUtils.getResponseBody(response, true);
167+
// The explain output should contain the highlight clause
168+
assertTrue("Explain should contain highlight", body.contains("highlight"));
169+
}
170+
171+
private JSONObject executeQueryWithHighlight(String query, String highlightJson)
172+
throws IOException {
173+
Request request = new Request("POST", "/_plugins/_ppl");
174+
request.setJsonEntity(
175+
String.format(Locale.ROOT, "{\"query\": \"%s\", \"highlight\": %s}", query, highlightJson));
176+
RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
177+
restOptionsBuilder.addHeader("Content-Type", "application/json");
178+
request.setOptions(restOptionsBuilder);
179+
Response response = client().performRequest(request);
180+
assertEquals(200, response.getStatusLine().getStatusCode());
181+
String body = org.opensearch.sql.legacy.TestUtils.getResponseBody(response, true);
182+
return new JSONObject(body);
183+
}
184+
}

0 commit comments

Comments
 (0)