Skip to content

Commit de3df61

Browse files
nathaliellenaanathaliellenaa
authored
Delete Memory Container API (opensearch-project#4027)
* Add delete memory container API Signed-off-by: Nathalie Jonathan <[email protected]> * Address comments Signed-off-by: Nathalie Jonathan <[email protected]> * Address comments, add UTs Signed-off-by: Nathalie Jonathan <[email protected]> * Add license header Signed-off-by: Nathalie Jonathan <[email protected]> * Apply spotless Signed-off-by: Nathalie Jonathan <[email protected]> * Refactor memory access control helper Signed-off-by: Nathalie Jonathan <[email protected]> * clean up Signed-off-by: Nathalie Jonathan <[email protected]> * Delete unused import Signed-off-by: Nathalie Jonathan <[email protected]> * Add more UTs Signed-off-by: Nathalie Jonathan <[email protected]> * Refactor Signed-off-by: Nathalie Jonathan <[email protected]> --------- Signed-off-by: Nathalie Jonathan <[email protected]>
1 parent 02601dc commit de3df61

File tree

8 files changed

+940
-0
lines changed

8 files changed

+940
-0
lines changed

‎common/src/main/java/org/opensearch/ml/common/transport/memorycontainer/MLMemoryContainerDeleteAction.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.memorycontainer;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.action.delete.DeleteResponse;
10+
11+
public class MLMemoryContainerDeleteAction extends ActionType<DeleteResponse> {
12+
public static final MLMemoryContainerDeleteAction INSTANCE = new MLMemoryContainerDeleteAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/memory_containers/delete";
14+
15+
private MLMemoryContainerDeleteAction() {
16+
super(NAME, DeleteResponse::new);
17+
}
18+
}

‎common/src/main/java/org/opensearch/ml/common/transport/memorycontainer/MLMemoryContainerDeleteRequest.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.memorycontainer;
7+
8+
import static org.opensearch.action.ValidateActions.addValidationError;
9+
10+
import java.io.ByteArrayInputStream;
11+
import java.io.ByteArrayOutputStream;
12+
import java.io.IOException;
13+
import java.io.UncheckedIOException;
14+
15+
import org.opensearch.action.ActionRequest;
16+
import org.opensearch.action.ActionRequestValidationException;
17+
import org.opensearch.core.common.io.stream.InputStreamStreamInput;
18+
import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
19+
import org.opensearch.core.common.io.stream.StreamInput;
20+
import org.opensearch.core.common.io.stream.StreamOutput;
21+
22+
import lombok.Builder;
23+
import lombok.Getter;
24+
25+
public class MLMemoryContainerDeleteRequest extends ActionRequest {
26+
@Getter
27+
String memoryContainerId;
28+
29+
@Getter
30+
String tenantId;
31+
32+
@Builder
33+
public MLMemoryContainerDeleteRequest(String memoryContainerId, String tenantId) {
34+
this.memoryContainerId = memoryContainerId;
35+
this.tenantId = tenantId;
36+
}
37+
38+
public MLMemoryContainerDeleteRequest(StreamInput input) throws IOException {
39+
super(input);
40+
this.memoryContainerId = input.readString();
41+
this.tenantId = input.readOptionalString();
42+
}
43+
44+
@Override
45+
public void writeTo(StreamOutput output) throws IOException {
46+
super.writeTo(output);
47+
output.writeString(memoryContainerId);
48+
output.writeOptionalString(tenantId);
49+
}
50+
51+
@Override
52+
public ActionRequestValidationException validate() {
53+
ActionRequestValidationException exception = null;
54+
55+
if (this.memoryContainerId == null) {
56+
exception = addValidationError("ML memory container id can't be null", exception);
57+
}
58+
59+
return exception;
60+
}
61+
62+
public static MLMemoryContainerDeleteRequest fromActionRequest(ActionRequest actionRequest) {
63+
if (actionRequest instanceof MLMemoryContainerDeleteRequest) {
64+
return (MLMemoryContainerDeleteRequest) actionRequest;
65+
}
66+
67+
try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
68+
actionRequest.writeTo(osso);
69+
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
70+
return new MLMemoryContainerDeleteRequest(input);
71+
}
72+
} catch (IOException e) {
73+
throw new UncheckedIOException("failed to parse ActionRequest into MLMemoryContainerDeleteRequest", e);
74+
}
75+
}
76+
}

‎common/src/test/java/org/opensearch/ml/common/transport/memorycontainer/MLMemoryContainerDeleteRequestTest.java‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.memorycontainer;
7+
8+
import static org.junit.Assert.assertEquals;
9+
import static org.junit.Assert.assertNotSame;
10+
import static org.junit.Assert.assertNull;
11+
import static org.junit.Assert.assertSame;
12+
13+
import java.io.IOException;
14+
import java.io.UncheckedIOException;
15+
16+
import org.junit.Before;
17+
import org.junit.Test;
18+
import org.opensearch.action.ActionRequest;
19+
import org.opensearch.action.ActionRequestValidationException;
20+
import org.opensearch.common.io.stream.BytesStreamOutput;
21+
import org.opensearch.core.common.io.stream.StreamOutput;
22+
23+
public class MLMemoryContainerDeleteRequestTest {
24+
private String memoryContainerId;
25+
26+
@Before
27+
public void setUp() {
28+
memoryContainerId = "test_id";
29+
}
30+
31+
@Test
32+
public void writeTo_Success() throws IOException {
33+
MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequest = MLMemoryContainerDeleteRequest
34+
.builder()
35+
.memoryContainerId(memoryContainerId)
36+
.build();
37+
BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
38+
mlMemoryContainerDeleteRequest.writeTo(bytesStreamOutput);
39+
MLMemoryContainerDeleteRequest parsedMemoryContainer = new MLMemoryContainerDeleteRequest(bytesStreamOutput.bytes().streamInput());
40+
assertEquals(parsedMemoryContainer.getMemoryContainerId(), memoryContainerId);
41+
}
42+
43+
@Test
44+
public void validate_Success() {
45+
MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequest = MLMemoryContainerDeleteRequest
46+
.builder()
47+
.memoryContainerId(memoryContainerId)
48+
.build();
49+
ActionRequestValidationException actionRequestValidationException = mlMemoryContainerDeleteRequest.validate();
50+
assertNull(actionRequestValidationException);
51+
}
52+
53+
@Test
54+
public void validate_Exception_NullMemoryContainerId() {
55+
MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequest = MLMemoryContainerDeleteRequest.builder().build();
56+
57+
ActionRequestValidationException exception = mlMemoryContainerDeleteRequest.validate();
58+
assertEquals("Validation Failed: 1: ML memory container id can't be null;", exception.getMessage());
59+
}
60+
61+
@Test
62+
public void fromActionRequest_Success() {
63+
MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequest = MLMemoryContainerDeleteRequest
64+
.builder()
65+
.memoryContainerId(memoryContainerId)
66+
.build();
67+
ActionRequest actionRequest = new ActionRequest() {
68+
@Override
69+
public ActionRequestValidationException validate() {
70+
return null;
71+
}
72+
73+
@Override
74+
public void writeTo(StreamOutput out) throws IOException {
75+
mlMemoryContainerDeleteRequest.writeTo(out);
76+
}
77+
};
78+
MLMemoryContainerDeleteRequest result = MLMemoryContainerDeleteRequest.fromActionRequest(actionRequest);
79+
assertNotSame(result, mlMemoryContainerDeleteRequest);
80+
assertEquals(result.getMemoryContainerId(), mlMemoryContainerDeleteRequest.getMemoryContainerId());
81+
}
82+
83+
@Test(expected = UncheckedIOException.class)
84+
public void fromActionRequest_IOException() {
85+
ActionRequest actionRequest = new ActionRequest() {
86+
@Override
87+
public ActionRequestValidationException validate() {
88+
return null;
89+
}
90+
91+
@Override
92+
public void writeTo(StreamOutput out) throws IOException {
93+
throw new IOException("test");
94+
}
95+
};
96+
MLMemoryContainerDeleteRequest.fromActionRequest(actionRequest);
97+
}
98+
99+
@Test
100+
public void fromActionRequestWithMemoryContainerDeleteRequest_Success() {
101+
MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequest = MLMemoryContainerDeleteRequest
102+
.builder()
103+
.memoryContainerId(memoryContainerId)
104+
.build();
105+
MLMemoryContainerDeleteRequest mlMemoryContainerDeleteRequestFromActionRequest = MLMemoryContainerDeleteRequest
106+
.fromActionRequest(mlMemoryContainerDeleteRequest);
107+
assertSame(mlMemoryContainerDeleteRequest, mlMemoryContainerDeleteRequestFromActionRequest);
108+
assertEquals(
109+
mlMemoryContainerDeleteRequest.getMemoryContainerId(),
110+
mlMemoryContainerDeleteRequestFromActionRequest.getMemoryContainerId()
111+
);
112+
}
113+
114+
@Test
115+
public void writeTo_withTenantId_Success() throws IOException {
116+
String tenantId = "tenant-1";
117+
MLMemoryContainerDeleteRequest request = MLMemoryContainerDeleteRequest
118+
.builder()
119+
.memoryContainerId(memoryContainerId)
120+
.tenantId(tenantId)
121+
.build();
122+
123+
BytesStreamOutput out = new BytesStreamOutput();
124+
request.writeTo(out);
125+
MLMemoryContainerDeleteRequest parsedRequest = new MLMemoryContainerDeleteRequest(out.bytes().streamInput());
126+
127+
assertEquals(memoryContainerId, parsedRequest.getMemoryContainerId());
128+
assertEquals(tenantId, parsedRequest.getTenantId());
129+
}
130+
131+
@Test
132+
public void writeTo_withoutTenantId_Success() throws IOException {
133+
MLMemoryContainerDeleteRequest request = MLMemoryContainerDeleteRequest.builder().memoryContainerId(memoryContainerId).build();
134+
135+
BytesStreamOutput out = new BytesStreamOutput();
136+
request.writeTo(out);
137+
MLMemoryContainerDeleteRequest parsedRequest = new MLMemoryContainerDeleteRequest(out.bytes().streamInput());
138+
139+
assertEquals(memoryContainerId, parsedRequest.getMemoryContainerId());
140+
assertNull(parsedRequest.getTenantId());
141+
}
142+
143+
@Test
144+
public void fromActionRequest_withTenantId_Success() {
145+
MLMemoryContainerDeleteRequest originalRequest = MLMemoryContainerDeleteRequest
146+
.builder()
147+
.memoryContainerId(memoryContainerId)
148+
.tenantId("tenant-1")
149+
.build();
150+
ActionRequest actionRequest = new ActionRequest() {
151+
@Override
152+
public ActionRequestValidationException validate() {
153+
return null;
154+
}
155+
156+
@Override
157+
public void writeTo(StreamOutput out) throws IOException {
158+
originalRequest.writeTo(out);
159+
}
160+
};
161+
162+
MLMemoryContainerDeleteRequest parsedRequest = MLMemoryContainerDeleteRequest.fromActionRequest(actionRequest);
163+
assertNotSame(originalRequest, parsedRequest);
164+
assertEquals(originalRequest.getMemoryContainerId(), parsedRequest.getMemoryContainerId());
165+
assertEquals(originalRequest.getTenantId(), parsedRequest.getTenantId());
166+
}
167+
}

0 commit comments

Comments
 (0)