Skip to content

Commit 9a48b76

Browse files
zane-neozane-neo
committed
Add UT file for testing file to check on the jacoco result
Signed-off-by: zane-neo <zaniu@amazon.com>
1 parent e1ee726 commit 9a48b76

File tree

1 file changed

+177
-0
lines changed

1 file changed

+177
-0
lines changed

‎plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsForTestingTests.java

+177
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.utils;
7+
8+
import com.fasterxml.jackson.core.JsonParseException;
9+
import org.junit.Assert;
10+
import org.junit.Test;
11+
import org.opensearch.Version;
12+
import org.opensearch.cluster.node.DiscoveryNode;
13+
import org.opensearch.cluster.node.DiscoveryNodeRole;
14+
import org.opensearch.common.xcontent.XContentFactory;
15+
import org.opensearch.core.common.bytes.BytesReference;
16+
import org.opensearch.core.xcontent.NamedXContentRegistry;
17+
import org.opensearch.core.xcontent.ToXContent;
18+
import org.opensearch.core.xcontent.XContentBuilder;
19+
import org.opensearch.core.xcontent.XContentParser;
20+
import org.opensearch.ml.common.MLTask;
21+
import org.opensearch.test.OpenSearchTestCase;
22+
23+
import java.io.IOException;
24+
import java.util.HashSet;
25+
import java.util.Set;
26+
27+
import static java.util.Collections.emptyMap;
28+
import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE;
29+
import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE;
30+
import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE;
31+
import static org.opensearch.ml.utils.TestHelper.ML_ROLE;
32+
33+
public class MLNodeUtilsForTestingTests extends OpenSearchTestCase {
34+
35+
public void testIsMLNode() {
36+
Set<DiscoveryNodeRole> roleSet = new HashSet<>();
37+
roleSet.add(DiscoveryNodeRole.DATA_ROLE);
38+
roleSet.add(DiscoveryNodeRole.INGEST_ROLE);
39+
DiscoveryNode normalNode = new DiscoveryNode("Normal node", buildNewFakeTransportAddress(), emptyMap(), roleSet, Version.CURRENT);
40+
Assert.assertFalse(MLNodeUtilsForTesting.isMLNode(normalNode));
41+
42+
roleSet.add(ML_ROLE);
43+
DiscoveryNode mlNode = new DiscoveryNode("ML node", buildNewFakeTransportAddress(), emptyMap(), roleSet, Version.CURRENT);
44+
Assert.assertTrue(MLNodeUtilsForTesting.isMLNode(mlNode));
45+
}
46+
47+
public void testCreateXContentParserFromRegistry() throws IOException {
48+
MLTask mlTask = MLTask.builder().taskId("taskId").modelId("modelId").build();
49+
XContentBuilder content = mlTask.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS);
50+
BytesReference bytesReference = BytesReference.bytes(content);
51+
NamedXContentRegistry namedXContentRegistry = NamedXContentRegistry.EMPTY;
52+
XContentParser xContentParser = MLNodeUtilsForTesting.createXContentParserFromRegistry(namedXContentRegistry, bytesReference);
53+
xContentParser.nextToken();
54+
MLTask parsedMLTask = MLTask.parse(xContentParser);
55+
assertEquals(mlTask, parsedMLTask);
56+
}
57+
58+
@Test
59+
public void testValidateSchema() throws IOException {
60+
String schema = "{"
61+
+ "\"type\": \"object\","
62+
+ "\"properties\": {"
63+
+ " \"key1\": {\"type\": \"string\"},"
64+
+ " \"key2\": {\"type\": \"integer\"}"
65+
+ "}"
66+
+ "}";
67+
String json = "{\"key1\": \"foo\", \"key2\": 123}";
68+
MLNodeUtilsForTesting.validateSchema(schema, json);
69+
}
70+
71+
@Test
72+
public void testValidateEmbeddingInputWithGeneralEmbeddingRemoteSchema() throws IOException {
73+
String schema = BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE.get("input");
74+
String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
75+
MLNodeUtilsForTesting.validateSchema(schema, json);
76+
}
77+
78+
@Test
79+
public void testValidateRemoteInputWithGeneralEmbeddingRemoteSchema() throws IOException {
80+
String schema = BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE.get("input");
81+
String json = "{\"parameters\": {\"texts\": [\"Hello\",\"world\"]}}";
82+
MLNodeUtilsForTesting.validateSchema(schema, json);
83+
}
84+
85+
@Test
86+
public void testValidateEmbeddingInputWithTitanTextRemoteSchema() throws IOException {
87+
String schema = BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE.get("input");
88+
String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
89+
MLNodeUtilsForTesting.validateSchema(schema, json);
90+
}
91+
92+
@Test
93+
public void testValidateRemoteInputWithTitanTextRemoteSchema() throws IOException {
94+
String schema = BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE.get("input");
95+
String json = "{\"parameters\": {\"inputText\": \"Say this is a test\"}}";
96+
MLNodeUtilsForTesting.validateSchema(schema, json);
97+
}
98+
99+
@Test
100+
public void testValidateEmbeddingInputWithTitanMultiModalRemoteSchema() throws IOException {
101+
String schema = BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE.get("input");
102+
String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
103+
MLNodeUtilsForTesting.validateSchema(schema, json);
104+
}
105+
106+
@Test
107+
public void testValidateRemoteInputWithTitanMultiModalRemoteSchema() throws IOException {
108+
String schema = BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE.get("input");
109+
String json = "{\n"
110+
+ " \"parameters\": {\n"
111+
+ " \"inputText\": \"Say this is a test\",\n"
112+
+ " \"inputImage\": \"/9jk=\"\n"
113+
+ " }\n"
114+
+ "}";
115+
MLNodeUtilsForTesting.validateSchema(schema, json);
116+
}
117+
118+
@Test
119+
public void testProcessRemoteInferenceInputDataSetParametersValueNoParameters() throws IOException {
120+
String json = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true}";
121+
String processedJson = MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json);
122+
assertEquals(json, processedJson);
123+
}
124+
125+
@Test
126+
public void testProcessRemoteInferenceInputDataSetInvalidJson() {
127+
String json = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{\"a\"}}";
128+
assertThrows(JsonParseException.class, () -> MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json));
129+
}
130+
131+
@Test
132+
public void testProcessRemoteInferenceInputDataSetEmptyParameters() throws IOException {
133+
String json = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{}}";
134+
String processedJson = MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json);
135+
assertEquals(json, processedJson);
136+
}
137+
138+
@Test
139+
public void testProcessRemoteInferenceInputDataSetParametersValueParametersWrongType() throws IOException {
140+
String json = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":[\"Hello\",\"world\"]}";
141+
String processedJson = MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json);
142+
assertEquals(json, processedJson);
143+
}
144+
145+
@Test
146+
public void testProcessRemoteInferenceInputDataSetParametersValueWithParametersProcessArray() throws IOException {
147+
String json = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{\"texts\":\"[\\\"Hello\\\",\\\"world\\\"]\"}}";
148+
String expectedJson = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{\"texts\":[\"Hello\",\"world\"]}}";
149+
String processedJson = MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json);
150+
assertEquals(expectedJson, processedJson);
151+
}
152+
153+
@Test
154+
public void testProcessRemoteInferenceInputDataSetParametersValueWithParametersProcessObject() throws IOException {
155+
String json =
156+
"{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{\"messages\":\"{\\\"role\\\":\\\"system\\\",\\\"foo\\\":\\\"{\\\\\\\"a\\\\\\\": \\\\\\\"b\\\\\\\"}\\\",\\\"content\\\":{\\\"a\\\":\\\"b\\\"}}\"}}}";
157+
String expectedJson =
158+
"{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{\"messages\":{\"role\":\"system\",\"foo\":\"{\\\"a\\\": \\\"b\\\"}\",\"content\":{\"a\":\"b\"}}}}";
159+
String processedJson = MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json);
160+
assertEquals(expectedJson, processedJson);
161+
}
162+
163+
@Test
164+
public void testProcessRemoteInferenceInputDataSetParametersValueWithParametersNoProcess() throws IOException {
165+
String json = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{\"key1\":\"foo\",\"key2\":123,\"key3\":true}}";
166+
String processedJson = MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json);
167+
assertEquals(json, processedJson);
168+
}
169+
170+
@Test
171+
public void testProcessRemoteInferenceInputDataSetParametersValueWithParametersInvalidJson() throws IOException {
172+
String json =
173+
"{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"parameters\":{\"key1\":\"foo\",\"key2\":123,\"key3\":true,\"texts\":\"[\\\"Hello\\\",\\\"world\\\"\"}}";
174+
String processedJson = MLNodeUtilsForTesting.processRemoteInferenceInputDataSetParametersValue(json);
175+
assertEquals(json, processedJson);
176+
}
177+
}

0 commit comments

Comments
 (0)