Skip to content

Commit 78eb74b

Browse files
Fix JsonGenerationException error in Local Sample Calculator and Anomaly Localization Execution Response (opensearch-project#3434)
* Add root object wrapper in LocalSampleCalculatorOutput and AnomalyLocalizationOutput. Signed-off-by: Nathalie Jonathan <[email protected]> * Fix format violations. Signed-off-by: Nathalie Jonathan <[email protected]> * Modify import * in RestMLExecuteActionTests, remove root object wrapper in AnomalyLocalizationOutputTests and LocalSampleCalculatorOutputTest. Signed-off-by: Nathalie Jonathan <[email protected]> * Removed unused variable in testAnomalyLocalizationExecutionResponse function. Signed-off-by: Nathalie Jonathan <[email protected]> --------- Signed-off-by: Nathalie Jonathan <[email protected]>
1 parent 17251cd commit 78eb74b

File tree

6 files changed

+113
-8
lines changed

6 files changed

+113
-8
lines changed

common/src/main/java/org/opensearch/ml/common/output/execute/anomalylocalization/AnomalyLocalizationOutput.java

+2
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,7 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
188188
@Override
189189
@SneakyThrows
190190
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) {
191+
builder.startObject();
191192
builder.startArray(FIELD_RESULTS);
192193
for (Map.Entry<String, Result> entry : this.results.entrySet()) {
193194
builder.startObject();
@@ -196,6 +197,7 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
196197
builder.endObject();
197198
}
198199
builder.endArray();
200+
builder.endObject();
199201
return builder;
200202
}
201203

common/src/main/java/org/opensearch/ml/common/output/execute/samplecalculator/LocalSampleCalculatorOutput.java

+2
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,11 @@ public void writeTo(StreamOutput out) throws IOException {
3939

4040
@Override
4141
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
42+
builder.startObject();
4243
if (result != null) {
4344
builder.field("result", result);
4445
}
46+
builder.endObject();
4547
return builder;
4648
}
4749
}

common/src/test/java/org/opensearch/ml/common/output/execute/anomalylocalization/AnomalyLocalizationOutputTests.java

-2
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,7 @@ public void testWriteable() throws Exception {
4747
@Test
4848
public void testXContent() throws Exception {
4949
XContentBuilder builder = XContentFactory.jsonBuilder();
50-
builder.startObject();
5150
builder = output.toXContent(builder, null);
52-
builder.endObject();
5351
String json = builder.toString();
5452
XContentParser parser = XContentType.JSON.xContent().createParser(NamedXContentRegistry.EMPTY, null, json);
5553
AnomalyLocalizationOutput newOutput = AnomalyLocalizationOutput.parse(parser);

common/src/test/java/org/opensearch/ml/common/output/execute/samplecalculator/LocalSampleCalculatorOutputTest.java

+2-6
Original file line numberDiff line numberDiff line change
@@ -28,21 +28,17 @@ public void setUp() {
2828

2929
@Test
3030
public void toXContent() throws IOException {
31-
XContentBuilder builder = XContentFactory.jsonBuilder();
32-
builder.startObject();
31+
XContentBuilder builder = MediaTypeRegistry.contentBuilder(XContentType.JSON);
3332
output.toXContent(builder, ToXContent.EMPTY_PARAMS);
34-
builder.endObject();
3533
String jsonStr = builder.toString();
3634
assertEquals("{\"result\":1.0}", jsonStr);
3735
}
3836

3937
@Test
4038
public void toXContent_EmptyOutput() throws IOException {
4139
LocalSampleCalculatorOutput output = LocalSampleCalculatorOutput.builder().build();
42-
XContentBuilder builder = XContentFactory.jsonBuilder();
43-
builder.startObject();
40+
XContentBuilder builder = MediaTypeRegistry.contentBuilder(XContentType.JSON);
4441
output.toXContent(builder, ToXContent.EMPTY_PARAMS);
45-
builder.endObject();
4642
String jsonStr = builder.toString();
4743
assertEquals("{}", jsonStr);
4844
}

plugin/src/test/java/org/opensearch/ml/rest/RestMLExecuteActionTests.java

+84
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,16 @@
1515
import static org.mockito.Mockito.times;
1616
import static org.mockito.Mockito.verify;
1717
import static org.mockito.Mockito.when;
18+
import static org.opensearch.ml.utils.TestHelper.getAnomalyLocalizationRestRequest;
1819
import static org.opensearch.ml.utils.TestHelper.getExecuteAgentRestRequest;
1920
import static org.opensearch.ml.utils.TestHelper.getLocalSampleCalculatorRestRequest;
2021
import static org.opensearch.ml.utils.TestHelper.getMetricsCorrelationRestRequest;
2122

2223
import java.io.IOException;
24+
import java.util.Arrays;
25+
import java.util.HashMap;
2326
import java.util.List;
27+
import java.util.Map;
2428

2529
import org.junit.Before;
2630
import org.mockito.ArgumentCaptor;
@@ -32,8 +36,13 @@
3236
import org.opensearch.core.action.ActionListener;
3337
import org.opensearch.core.common.Strings;
3438
import org.opensearch.core.rest.RestStatus;
39+
import org.opensearch.core.xcontent.XContentBuilder;
3540
import org.opensearch.ml.common.FunctionName;
3641
import org.opensearch.ml.common.input.Input;
42+
import org.opensearch.ml.common.output.execute.anomalylocalization.AnomalyLocalizationOutput;
43+
import org.opensearch.ml.common.output.execute.anomalylocalization.AnomalyLocalizationOutput.Bucket;
44+
import org.opensearch.ml.common.output.execute.anomalylocalization.AnomalyLocalizationOutput.Result;
45+
import org.opensearch.ml.common.output.execute.samplecalculator.LocalSampleCalculatorOutput;
3746
import org.opensearch.ml.common.transport.execute.MLExecuteTaskAction;
3847
import org.opensearch.ml.common.transport.execute.MLExecuteTaskRequest;
3948
import org.opensearch.ml.common.transport.execute.MLExecuteTaskResponse;
@@ -337,4 +346,79 @@ public void testAgentExecutionResponsePlainText() throws Exception {
337346
"{\"error\":{\"reason\":\"Invalid Request\",\"details\":\"Illegal Argument Exception\",\"type\":\"IllegalArgumentException\"},\"status\":400}";
338347
assertEquals(expectedError, response.content().utf8ToString());
339348
}
349+
350+
public void testLocalSampleCalculatorExecutionResponse() throws Exception {
351+
RestRequest request = getLocalSampleCalculatorRestRequest();
352+
XContentBuilder builder = XContentFactory.jsonBuilder();
353+
when(channel.newBuilder()).thenReturn(builder);
354+
doAnswer(invocation -> {
355+
ActionListener<MLExecuteTaskResponse> actionListener = invocation.getArgument(2);
356+
LocalSampleCalculatorOutput output = LocalSampleCalculatorOutput.builder().totalSum(3.0).build();
357+
MLExecuteTaskResponse response = MLExecuteTaskResponse
358+
.builder()
359+
.output(output)
360+
.functionName(FunctionName.LOCAL_SAMPLE_CALCULATOR)
361+
.build();
362+
actionListener.onResponse(response);
363+
return null;
364+
}).when(client).execute(eq(MLExecuteTaskAction.INSTANCE), any(), any());
365+
doNothing().when(channel).sendResponse(any());
366+
restMLExecuteAction.handleRequest(request, channel, client);
367+
368+
ArgumentCaptor<RestResponse> responseCaptor = ArgumentCaptor.forClass(RestResponse.class);
369+
verify(channel).sendResponse(responseCaptor.capture());
370+
BytesRestResponse response = (BytesRestResponse) responseCaptor.getValue();
371+
assertEquals(RestStatus.OK, response.status());
372+
assertEquals("{\"result\":3.0}", response.content().utf8ToString());
373+
}
374+
375+
public void testAnomalyLocalizationExecutionResponse() throws Exception {
376+
RestRequest request = getAnomalyLocalizationRestRequest();
377+
XContentBuilder builder = XContentFactory.jsonBuilder();
378+
when(channel.newBuilder()).thenReturn(builder);
379+
doAnswer(invocation -> {
380+
ActionListener<MLExecuteTaskResponse> actionListener = invocation.getArgument(2);
381+
382+
Bucket bucket1 = new Bucket();
383+
bucket1.setStartTime(1620630000000L);
384+
bucket1.setEndTime(1620716400000L);
385+
bucket1.setOverallAggValue(65.0);
386+
387+
Result result = new Result();
388+
result.setBuckets(Arrays.asList(bucket1));
389+
390+
AnomalyLocalizationOutput output = new AnomalyLocalizationOutput();
391+
Map<String, Result> results = new HashMap<>();
392+
results.put("sum", result);
393+
output.setResults(results);
394+
395+
MLExecuteTaskResponse response = MLExecuteTaskResponse
396+
.builder()
397+
.output(output)
398+
.functionName(FunctionName.ANOMALY_LOCALIZATION)
399+
.build();
400+
actionListener.onResponse(response);
401+
return null;
402+
}).when(client).execute(eq(MLExecuteTaskAction.INSTANCE), any(), any());
403+
doNothing().when(channel).sendResponse(any());
404+
restMLExecuteAction.handleRequest(request, channel, client);
405+
406+
ArgumentCaptor<RestResponse> responseCaptor = ArgumentCaptor.forClass(RestResponse.class);
407+
verify(channel).sendResponse(responseCaptor.capture());
408+
BytesRestResponse response = (BytesRestResponse) responseCaptor.getValue();
409+
assertEquals(RestStatus.OK, response.status());
410+
String expectedJson = "{\"results\":[{"
411+
+ "\"name\":\"sum\","
412+
+ "\"result\":{"
413+
+ "\"buckets\":["
414+
+ "{"
415+
+ "\"start_time\":1620630000000,"
416+
+ "\"end_time\":1620716400000,"
417+
+ "\"overall_aggregate_value\":65.0"
418+
+ "}"
419+
+ "]"
420+
+ "}"
421+
+ "}]}";
422+
assertEquals(expectedJson, response.content().utf8ToString());
423+
}
340424
}

plugin/src/test/java/org/opensearch/ml/utils/TestHelper.java

+23
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@
8080
import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
8181
import org.opensearch.ml.common.input.Constants;
8282
import org.opensearch.ml.common.input.MLInput;
83+
import org.opensearch.ml.common.input.execute.anomalylocalization.AnomalyLocalizationInput;
8384
import org.opensearch.ml.common.input.execute.metricscorrelation.MetricsCorrelationInput;
8485
import org.opensearch.ml.common.input.execute.samplecalculator.LocalSampleCalculatorInput;
8586
import org.opensearch.ml.common.input.parameter.clustering.KMeansParams;
@@ -360,6 +361,27 @@ public static RestRequest getMetricsCorrelationRestRequest() {
360361
.build();
361362
}
362363

364+
public static RestRequest getAnomalyLocalizationRestRequest() {
365+
Map<String, String> params = new HashMap<>();
366+
params.put(PARAMETER_ALGORITHM, FunctionName.ANOMALY_LOCALIZATION.name());
367+
final String requestContent = "{"
368+
+ "\"input_data\": {"
369+
+ "\"index_name\": \"test-index\","
370+
+ "\"attribute_field_names\": [\"attribute\"],"
371+
+ "\"time_field_name\": \"timestamp\","
372+
+ "\"start_time\": 1620630000000,"
373+
+ "\"end_time\": 1621234800000,"
374+
+ "\"min_time_interval\": 86400000,"
375+
+ "\"num_outputs\": 1"
376+
+ "}"
377+
+ "}";
378+
RestRequest request = new FakeRestRequest.Builder(getXContentRegistry())
379+
.withParams(params)
380+
.withContent(new BytesArray(requestContent), XContentType.JSON)
381+
.build();
382+
return request;
383+
}
384+
363385
public static RestRequest getExecuteAgentRestRequest() {
364386
Map<String, String> params = new HashMap<>();
365387
params.put(PARAMETER_AGENT_ID, "test_agent_id");
@@ -407,6 +429,7 @@ private static NamedXContentRegistry getXContentRegistry() {
407429
entries.add(KMeansParams.XCONTENT_REGISTRY);
408430
entries.add(LocalSampleCalculatorInput.XCONTENT_REGISTRY);
409431
entries.add(MetricsCorrelationInput.XCONTENT_REGISTRY);
432+
entries.add(AnomalyLocalizationInput.XCONTENT_REGISTRY_ENTRY);
410433
return new NamedXContentRegistry(entries);
411434
}
412435

0 commit comments

Comments
 (0)