From 02b0d5626763d761bc8dfada605658b0fefa4345 Mon Sep 17 00:00:00 2001 From: Bhavana Ramaram Date: Tue, 3 Sep 2024 18:51:21 -0500 Subject: [PATCH] fix breaking changes in config index fields Signed-off-by: Bhavana Ramaram --- .../org/opensearch/ml/common/MLConfig.java | 53 ++++++++++++++++--- .../config/MLConfigGetResponseTest.java | 4 +- .../config/GetConfigTransportActionTests.java | 25 ++++++++- 3 files changed, 73 insertions(+), 9 deletions(-) diff --git a/common/src/main/java/org/opensearch/ml/common/MLConfig.java b/common/src/main/java/org/opensearch/ml/common/MLConfig.java index ccdeed5df3..39dd250f95 100644 --- a/common/src/main/java/org/opensearch/ml/common/MLConfig.java +++ b/common/src/main/java/org/opensearch/ml/common/MLConfig.java @@ -42,48 +42,83 @@ public class MLConfig implements ToXContentObject, Writeable { @Setter private String type; + @Setter + private String configType; + private Configuration configuration; + private Configuration mlConfiguration; private final Instant createTime; private Instant lastUpdateTime; + private Instant lastUpdatedTime; @Builder(toBuilder = true) - public MLConfig(String type, Configuration configuration, Instant createTime, Instant lastUpdateTime) { + public MLConfig( + String type, + String configType, + Configuration configuration, + Configuration mlConfiguration, + Instant createTime, + Instant lastUpdateTime, + Instant lastUpdatedTime + ) { this.type = type; + this.configType = configType; this.configuration = configuration; + this.mlConfiguration = mlConfiguration; this.createTime = createTime; this.lastUpdateTime = lastUpdateTime; + this.lastUpdatedTime = lastUpdatedTime; } public MLConfig(StreamInput input) throws IOException { this.type = input.readOptionalString(); + this.configType = input.readOptionalString(); if (input.readBoolean()) { configuration = new Configuration(input); } + if (input.readBoolean()) { + mlConfiguration = new Configuration(input); + } createTime = input.readOptionalInstant(); lastUpdateTime = input.readOptionalInstant(); + lastUpdatedTime = input.readOptionalInstant(); } @Override public void writeTo(StreamOutput out) throws IOException { out.writeOptionalString(type); + out.writeOptionalString(configType); if (configuration != null) { out.writeBoolean(true); configuration.writeTo(out); } else { out.writeBoolean(false); } + if (mlConfiguration != null) { + out.writeBoolean(true); + mlConfiguration.writeTo(out); + } else { + out.writeBoolean(false); + } out.writeOptionalInstant(createTime); out.writeOptionalInstant(lastUpdateTime); + out.writeOptionalInstant(lastUpdatedTime); } @Override public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException { XContentBuilder builder = xContentBuilder.startObject(); if (type != null) { - builder.field(CONFIG_TYPE_FIELD, type); + builder.field(TYPE_FIELD, type); + } + if (configType != null) { + builder.field(CONFIG_TYPE_FIELD, configType); } if (configuration != null) { - builder.field(ML_CONFIGURATION_FIELD, configuration); + builder.field(CONFIGURATION_FIELD, configuration); + } + if (mlConfiguration != null) { + builder.field(ML_CONFIGURATION_FIELD, mlConfiguration); } if (createTime != null) { builder.field(CREATE_TIME_FIELD, createTime.toEpochMilli()); @@ -91,6 +126,9 @@ public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params if (lastUpdateTime != null) { builder.field(LAST_UPDATE_TIME_FIELD, lastUpdateTime.toEpochMilli()); } + if (lastUpdatedTime != null) { + builder.field(LAST_UPDATED_TIME_FIELD, lastUpdatedTime.toEpochMilli()); + } return builder.endObject(); } @@ -142,10 +180,13 @@ public static MLConfig parse(XContentParser parser) throws IOException { } return MLConfig .builder() - .type(configType == null ? type : configType) - .configuration(mlConfiguration == null ? configuration : mlConfiguration) + .type(type) + .configType(configType) + .configuration(configuration) + .mlConfiguration(mlConfiguration) .createTime(createTime) - .lastUpdateTime(lastUpdatedTime == null ? lastUpdateTime : lastUpdatedTime) + .lastUpdateTime(lastUpdateTime) + .lastUpdatedTime(lastUpdatedTime) .build(); } } diff --git a/common/src/test/java/org/opensearch/ml/common/transport/config/MLConfigGetResponseTest.java b/common/src/test/java/org/opensearch/ml/common/transport/config/MLConfigGetResponseTest.java index b187b4a8c8..5e24bcaa6e 100644 --- a/common/src/test/java/org/opensearch/ml/common/transport/config/MLConfigGetResponseTest.java +++ b/common/src/test/java/org/opensearch/ml/common/transport/config/MLConfigGetResponseTest.java @@ -60,7 +60,7 @@ public void MLConfigGetResponse_Builder() throws IOException { @Test public void writeTo() throws IOException { // create ml agent using mlConfig and mlConfigGetResponse - mlConfig = new MLConfig("olly_agent", new Configuration("agent_id"), Instant.EPOCH, Instant.EPOCH); + mlConfig = new MLConfig("olly_agent", null, new Configuration("agent_id"), null, Instant.EPOCH, Instant.EPOCH, Instant.EPOCH); MLConfigGetResponse mlConfigGetResponse = MLConfigGetResponse.builder().mlConfig(mlConfig).build(); // use write out for both agents BytesStreamOutput output = new BytesStreamOutput(); @@ -76,7 +76,7 @@ public void writeTo() throws IOException { @Test public void toXContent() throws IOException { - mlConfig = new MLConfig(null, null, null, null); + mlConfig = new MLConfig(null, null, null, null, null, null, null); MLConfigGetResponse mlConfigGetResponse = MLConfigGetResponse.builder().mlConfig(mlConfig).build(); XContentBuilder builder = XContentFactory.jsonBuilder(); ToXContent.Params params = EMPTY_PARAMS; diff --git a/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java index 0dbb79ef6c..afa4153a74 100644 --- a/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java +++ b/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java @@ -162,9 +162,32 @@ public void testDoExecute_Success() throws IOException { verify(actionListener).onResponse(any(MLConfigGetResponse.class)); } + @Test + public void testDoExecute_Success_ForNewFields() throws IOException { + String configID = "config_id"; + MLConfig mlConfig = new MLConfig(null, "olly_agent", null, new Configuration("agent_id"), Instant.EPOCH, null, Instant.EPOCH); + + XContentBuilder content = mlConfig.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS); + BytesReference bytesReference = BytesReference.bytes(content); + GetResult getResult = new GetResult("indexName", configID, 111l, 111l, 111l, true, bytesReference, null, null); + GetResponse getResponse = new GetResponse(getResult); + ActionListener actionListener = mock(ActionListener.class); + MLConfigGetRequest request = new MLConfigGetRequest(configID); + Task task = mock(Task.class); + + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(1); + listener.onResponse(getResponse); + return null; + }).when(client).get(any(), any()); + + getConfigTransportAction.doExecute(task, request, actionListener); + verify(actionListener).onResponse(any(MLConfigGetResponse.class)); + } + public GetResponse prepareMLConfig(String configID) throws IOException { - MLConfig mlConfig = new MLConfig("olly_agent", new Configuration("agent_id"), Instant.EPOCH, Instant.EPOCH); + MLConfig mlConfig = new MLConfig("olly_agent", null, new Configuration("agent_id"), null, Instant.EPOCH, Instant.EPOCH, null); XContentBuilder content = mlConfig.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS); BytesReference bytesReference = BytesReference.bytes(content);