Skip to content

Commit aed8983

Browse files
dhrubo-osdhrubo-os
authored
adding unit tests for create and get memory container functionalities (opensearch-project#4056)
Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>
1 parent 3fa3a1e commit aed8983

File tree

15 files changed

+5246
-16
lines changed

15 files changed

+5246
-16
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
@Getter
2626
public class MLCreateMemoryContainerRequest extends ActionRequest {
2727

28-
private MLCreateMemoryContainerInput mlCreateMemoryContainerInput;
28+
private final MLCreateMemoryContainerInput mlCreateMemoryContainerInput;
2929

3030
@Builder
3131
public MLCreateMemoryContainerRequest(MLCreateMemoryContainerInput mlCreateMemoryContainerInput) {

‎common/src/test/java/org/opensearch/ml/common/memorycontainer/MLMemoryContainerTests.java‎

Lines changed: 426 additions & 0 deletions
Large diffs are not rendered by default.

‎common/src/test/java/org/opensearch/ml/common/memorycontainer/MemoryStorageConfigTests.java‎

Lines changed: 413 additions & 0 deletions
Large diffs are not rendered by default.

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

Lines changed: 389 additions & 0 deletions
Large diffs are not rendered by default.

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

Lines changed: 304 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,304 @@
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.assertNotNull;
10+
import static org.junit.Assert.assertNull;
11+
import static org.junit.Assert.assertSame;
12+
import static org.junit.Assert.assertTrue;
13+
14+
import java.io.IOException;
15+
import java.io.UncheckedIOException;
16+
17+
import org.junit.Before;
18+
import org.junit.Test;
19+
import org.opensearch.action.ActionRequest;
20+
import org.opensearch.action.ActionRequestValidationException;
21+
import org.opensearch.common.io.stream.BytesStreamOutput;
22+
import org.opensearch.core.common.io.stream.StreamInput;
23+
import org.opensearch.core.common.io.stream.StreamOutput;
24+
import org.opensearch.ml.common.FunctionName;
25+
import org.opensearch.ml.common.memorycontainer.MemoryStorageConfig;
26+
27+
public class MLCreateMemoryContainerRequestTests {
28+
29+
private MLCreateMemoryContainerRequest requestWithAllFields;
30+
private MLCreateMemoryContainerRequest requestMinimal;
31+
private MLCreateMemoryContainerInput testInput;
32+
private MLCreateMemoryContainerInput minimalInput;
33+
private MemoryStorageConfig testMemoryStorageConfig;
34+
35+
@Before
36+
public void setUp() {
37+
// Create test memory storage config
38+
testMemoryStorageConfig = MemoryStorageConfig
39+
.builder()
40+
.memoryIndexName("test-memory-index")
41+
.embeddingModelType(FunctionName.TEXT_EMBEDDING)
42+
.embeddingModelId("test-embedding-model")
43+
.llmModelId("test-llm-model")
44+
.dimension(768)
45+
.maxInferSize(8)
46+
.build();
47+
48+
// Create test input with all fields
49+
testInput = MLCreateMemoryContainerInput
50+
.builder()
51+
.name("test-memory-container")
52+
.description("Test memory container description")
53+
.memoryStorageConfig(testMemoryStorageConfig)
54+
.tenantId("test-tenant")
55+
.build();
56+
57+
// Create minimal input
58+
minimalInput = MLCreateMemoryContainerInput.builder().name("minimal-container").build();
59+
60+
// Create requests
61+
requestWithAllFields = MLCreateMemoryContainerRequest.builder().mlCreateMemoryContainerInput(testInput).build();
62+
63+
requestMinimal = MLCreateMemoryContainerRequest.builder().mlCreateMemoryContainerInput(minimalInput).build();
64+
}
65+
66+
@Test
67+
public void testConstructorWithBuilder() {
68+
assertNotNull(requestWithAllFields);
69+
assertEquals(testInput, requestWithAllFields.getMlCreateMemoryContainerInput());
70+
71+
assertNotNull(requestMinimal);
72+
assertEquals(minimalInput, requestMinimal.getMlCreateMemoryContainerInput());
73+
}
74+
75+
@Test
76+
public void testConstructorWithInput() {
77+
MLCreateMemoryContainerRequest request = new MLCreateMemoryContainerRequest(testInput);
78+
79+
assertNotNull(request);
80+
assertEquals(testInput, request.getMlCreateMemoryContainerInput());
81+
}
82+
83+
@Test
84+
public void testStreamInputOutput() throws IOException {
85+
BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
86+
requestWithAllFields.writeTo(bytesStreamOutput);
87+
88+
StreamInput streamInput = bytesStreamOutput.bytes().streamInput();
89+
MLCreateMemoryContainerRequest parsedRequest = new MLCreateMemoryContainerRequest(streamInput);
90+
91+
assertNotNull(parsedRequest);
92+
assertNotNull(parsedRequest.getMlCreateMemoryContainerInput());
93+
94+
// Verify the input fields
95+
MLCreateMemoryContainerInput originalInput = requestWithAllFields.getMlCreateMemoryContainerInput();
96+
MLCreateMemoryContainerInput parsedInput = parsedRequest.getMlCreateMemoryContainerInput();
97+
98+
assertEquals(originalInput.getName(), parsedInput.getName());
99+
assertEquals(originalInput.getDescription(), parsedInput.getDescription());
100+
assertEquals(originalInput.getTenantId(), parsedInput.getTenantId());
101+
assertEquals(originalInput.getMemoryStorageConfig(), parsedInput.getMemoryStorageConfig());
102+
}
103+
104+
@Test
105+
public void testStreamInputOutputWithMinimalFields() throws IOException {
106+
BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
107+
requestMinimal.writeTo(bytesStreamOutput);
108+
109+
StreamInput streamInput = bytesStreamOutput.bytes().streamInput();
110+
MLCreateMemoryContainerRequest parsedRequest = new MLCreateMemoryContainerRequest(streamInput);
111+
112+
assertNotNull(parsedRequest);
113+
assertNotNull(parsedRequest.getMlCreateMemoryContainerInput());
114+
115+
MLCreateMemoryContainerInput parsedInput = parsedRequest.getMlCreateMemoryContainerInput();
116+
assertEquals("minimal-container", parsedInput.getName());
117+
assertNull(parsedInput.getDescription());
118+
assertNull(parsedInput.getTenantId());
119+
assertNull(parsedInput.getMemoryStorageConfig());
120+
}
121+
122+
@Test
123+
public void testValidateWithValidInput() {
124+
ActionRequestValidationException validationException = requestWithAllFields.validate();
125+
assertNull(validationException);
126+
}
127+
128+
@Test
129+
public void testValidateWithMinimalValidInput() {
130+
ActionRequestValidationException validationException = requestMinimal.validate();
131+
assertNull(validationException);
132+
}
133+
134+
@Test
135+
public void testValidateWithNullInput() {
136+
MLCreateMemoryContainerRequest requestWithNullInput = MLCreateMemoryContainerRequest
137+
.builder()
138+
.mlCreateMemoryContainerInput(null)
139+
.build();
140+
141+
ActionRequestValidationException validationException = requestWithNullInput.validate();
142+
143+
assertNotNull(validationException);
144+
assertTrue(validationException.validationErrors().contains("Memory container input can't be null"));
145+
}
146+
147+
@Test
148+
public void testFromActionRequestWithSameType() {
149+
MLCreateMemoryContainerRequest result = MLCreateMemoryContainerRequest.fromActionRequest(requestWithAllFields);
150+
151+
assertSame(requestWithAllFields, result);
152+
}
153+
154+
@Test
155+
public void testFromActionRequestWithDifferentType() throws IOException {
156+
// Create a properly serializable ActionRequest that writes data in the expected format
157+
ActionRequest mockActionRequest = new ActionRequest() {
158+
@Override
159+
public ActionRequestValidationException validate() {
160+
return null;
161+
}
162+
163+
@Override
164+
public void writeTo(StreamOutput out) throws IOException {
165+
// Write data in the same format as MLCreateMemoryContainerRequest
166+
super.writeTo(out); // Write ActionRequest base data
167+
testInput.writeTo(out); // Write the MLCreateMemoryContainerInput data
168+
}
169+
};
170+
171+
MLCreateMemoryContainerRequest result = MLCreateMemoryContainerRequest.fromActionRequest(mockActionRequest);
172+
173+
assertNotNull(result);
174+
assertNotNull(result.getMlCreateMemoryContainerInput());
175+
assertEquals(testInput.getName(), result.getMlCreateMemoryContainerInput().getName());
176+
assertEquals(testInput.getDescription(), result.getMlCreateMemoryContainerInput().getDescription());
177+
assertEquals(testInput.getTenantId(), result.getMlCreateMemoryContainerInput().getTenantId());
178+
}
179+
180+
@Test(expected = UncheckedIOException.class)
181+
public void testFromActionRequestWithIOException() {
182+
// Create a mock ActionRequest that throws IOException during serialization
183+
ActionRequest mockActionRequest = new ActionRequest() {
184+
@Override
185+
public ActionRequestValidationException validate() {
186+
return null;
187+
}
188+
189+
@Override
190+
public void writeTo(StreamOutput out) throws IOException {
191+
throw new IOException("Test IOException");
192+
}
193+
};
194+
195+
MLCreateMemoryContainerRequest.fromActionRequest(mockActionRequest);
196+
}
197+
198+
@Test
199+
public void testGetterFunctionality() {
200+
assertEquals(testInput, requestWithAllFields.getMlCreateMemoryContainerInput());
201+
assertEquals(minimalInput, requestMinimal.getMlCreateMemoryContainerInput());
202+
}
203+
204+
@Test
205+
public void testBuilderFunctionality() {
206+
MLCreateMemoryContainerRequest request = MLCreateMemoryContainerRequest.builder().mlCreateMemoryContainerInput(testInput).build();
207+
208+
assertNotNull(request);
209+
assertEquals(testInput, request.getMlCreateMemoryContainerInput());
210+
}
211+
212+
@Test
213+
public void testInheritanceFromActionRequest() {
214+
assertTrue(requestWithAllFields instanceof ActionRequest);
215+
assertTrue(requestMinimal instanceof ActionRequest);
216+
}
217+
218+
@Test
219+
public void testValidationDelegation() {
220+
// Test that validation is properly delegated to the input object
221+
// The request itself only validates that input is not null
222+
// All other validation is handled by MLCreateMemoryContainerInput and MemoryStorageConfig
223+
224+
// Valid input should pass validation
225+
ActionRequestValidationException validationException = requestWithAllFields.validate();
226+
assertNull(validationException);
227+
228+
// Null input should fail validation
229+
MLCreateMemoryContainerRequest nullInputRequest = MLCreateMemoryContainerRequest
230+
.builder()
231+
.mlCreateMemoryContainerInput(null)
232+
.build();
233+
234+
ActionRequestValidationException nullValidationException = nullInputRequest.validate();
235+
assertNotNull(nullValidationException);
236+
assertEquals(1, nullValidationException.validationErrors().size());
237+
assertTrue(nullValidationException.validationErrors().get(0).contains("Memory container input can't be null"));
238+
}
239+
240+
@Test
241+
public void testCompleteRoundTripSerialization() throws IOException {
242+
// Test complete serialization round trip
243+
BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
244+
requestWithAllFields.writeTo(bytesStreamOutput);
245+
246+
StreamInput streamInput = bytesStreamOutput.bytes().streamInput();
247+
MLCreateMemoryContainerRequest deserializedRequest = new MLCreateMemoryContainerRequest(streamInput);
248+
249+
// Verify all nested data is preserved
250+
MLCreateMemoryContainerInput originalInput = requestWithAllFields.getMlCreateMemoryContainerInput();
251+
MLCreateMemoryContainerInput deserializedInput = deserializedRequest.getMlCreateMemoryContainerInput();
252+
253+
assertEquals(originalInput.getName(), deserializedInput.getName());
254+
assertEquals(originalInput.getDescription(), deserializedInput.getDescription());
255+
assertEquals(originalInput.getTenantId(), deserializedInput.getTenantId());
256+
257+
// Verify nested MemoryStorageConfig
258+
MemoryStorageConfig originalConfig = originalInput.getMemoryStorageConfig();
259+
MemoryStorageConfig deserializedConfig = deserializedInput.getMemoryStorageConfig();
260+
261+
assertEquals(originalConfig.getMemoryIndexName(), deserializedConfig.getMemoryIndexName());
262+
assertEquals(originalConfig.isSemanticStorageEnabled(), deserializedConfig.isSemanticStorageEnabled());
263+
assertEquals(originalConfig.getEmbeddingModelType(), deserializedConfig.getEmbeddingModelType());
264+
assertEquals(originalConfig.getEmbeddingModelId(), deserializedConfig.getEmbeddingModelId());
265+
assertEquals(originalConfig.getLlmModelId(), deserializedConfig.getLlmModelId());
266+
assertEquals(originalConfig.getDimension(), deserializedConfig.getDimension());
267+
assertEquals(originalConfig.getMaxInferSize(), deserializedConfig.getMaxInferSize());
268+
}
269+
270+
@Test
271+
public void testFromActionRequestRoundTrip() throws IOException {
272+
// Test that fromActionRequest can properly handle the same request type
273+
MLCreateMemoryContainerRequest reconstructed = MLCreateMemoryContainerRequest.fromActionRequest(requestWithAllFields);
274+
assertSame(requestWithAllFields, reconstructed);
275+
276+
// Test with minimal request
277+
MLCreateMemoryContainerRequest minimalReconstructed = MLCreateMemoryContainerRequest.fromActionRequest(requestMinimal);
278+
assertSame(requestMinimal, minimalReconstructed);
279+
}
280+
281+
@Test
282+
public void testNullInputHandling() {
283+
MLCreateMemoryContainerRequest requestWithNull = MLCreateMemoryContainerRequest
284+
.builder()
285+
.mlCreateMemoryContainerInput(null)
286+
.build();
287+
288+
assertNotNull(requestWithNull);
289+
assertNull(requestWithNull.getMlCreateMemoryContainerInput());
290+
291+
// Validation should catch this
292+
ActionRequestValidationException validationException = requestWithNull.validate();
293+
assertNotNull(validationException);
294+
assertTrue(validationException.validationErrors().contains("Memory container input can't be null"));
295+
}
296+
297+
@Test
298+
public void testBuilderWithNullInput() {
299+
MLCreateMemoryContainerRequest request = MLCreateMemoryContainerRequest.builder().mlCreateMemoryContainerInput(null).build();
300+
301+
assertNotNull(request);
302+
assertNull(request.getMlCreateMemoryContainerInput());
303+
}
304+
}

0 commit comments

Comments
 (0)