Skip to content

Commit 43dcefc

Browse files
removing api keys from the integ test log (#3112) (#3124)
* removing api keys from the integ test log Signed-off-by: Dhrubo Saha <[email protected]> * apply spotless Signed-off-by: Dhrubo Saha <[email protected]> * addressed comment Signed-off-by: Dhrubo Saha <[email protected]> * adding unit test to address comments Signed-off-by: Dhrubo Saha <[email protected]> --------- Signed-off-by: Dhrubo Saha <[email protected]> (cherry picked from commit bc030fa) Co-authored-by: Dhrubo Saha <[email protected]>
1 parent 41dae61 commit 43dcefc

File tree

1 file changed

+57
-26
lines changed

1 file changed

+57
-26
lines changed

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

Lines changed: 57 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,10 @@
1717
import org.junit.Before;
1818
import org.junit.Ignore;
1919
import org.junit.Rule;
20+
import org.junit.Test;
2021
import org.junit.rules.ExpectedException;
2122
import org.opensearch.client.Response;
23+
import org.opensearch.client.ResponseException;
2224
import org.opensearch.ml.common.MLTaskState;
2325
import org.opensearch.ml.utils.TestHelper;
2426

@@ -76,39 +78,63 @@ public void setup() throws IOException, InterruptedException {
7678
Thread.sleep(20000);
7779
}
7880

79-
public void testCreateConnector() throws IOException {
80-
Response response = createConnector(completionModelConnectorEntity);
81-
Map responseMap = parseResponseToMap(response);
82-
assertNotNull((String) responseMap.get("connector_id"));
83-
}
84-
85-
public void testGetConnector() throws IOException {
81+
public void testCreate_Get_DeleteConnector() throws IOException {
8682
Response response = createConnector(completionModelConnectorEntity);
8783
Map responseMap = parseResponseToMap(response);
8884
String connectorId = (String) responseMap.get("connector_id");
85+
assertNotNull(connectorId); // Testing create connector
86+
87+
// Testing Get connector
8988
response = TestHelper.makeRequest(client(), "GET", "/_plugins/_ml/connectors/" + connectorId, null, "", null);
9089
responseMap = parseResponseToMap(response);
91-
assertEquals("OpenAI Connector", (String) responseMap.get("name"));
92-
assertEquals("1", (String) responseMap.get("version"));
93-
assertEquals("The connector to public OpenAI model service for GPT 3.5", (String) responseMap.get("description"));
94-
assertEquals("http", (String) responseMap.get("protocol"));
95-
}
90+
assertEquals("OpenAI Connector", responseMap.get("name"));
91+
assertEquals("1", responseMap.get("version"));
92+
assertEquals("The connector to public OpenAI model service for GPT 3.5", responseMap.get("description"));
93+
assertEquals("http", responseMap.get("protocol"));
9694

97-
public void testDeleteConnector() throws IOException {
98-
Response response = createConnector(completionModelConnectorEntity);
99-
Map responseMap = parseResponseToMap(response);
100-
String connectorId = (String) responseMap.get("connector_id");
95+
// Testing delete connector
10196
response = TestHelper.makeRequest(client(), "DELETE", "/_plugins/_ml/connectors/" + connectorId, null, "", null);
10297
responseMap = parseResponseToMap(response);
103-
assertEquals("deleted", (String) responseMap.get("result"));
98+
assertEquals("deleted", responseMap.get("result"));
99+
100+
}
101+
102+
private static String maskSensitiveInfo(String input) {
103+
// Regex to remove the whole credential object and replace it with "***"
104+
String regex = "\"credential\":\\{.*?}";
105+
return input.replaceAll(regex, "\"credential\": \"***\"");
106+
}
107+
108+
@Test
109+
public void testMaskSensitiveInfo_withCredential() {
110+
String input = "{\"credential\":{\"username\":\"admin\",\"password\":\"secret\"}}";
111+
String expectedOutput = "{\"credential\": \"***\"}";
112+
String actualOutput = maskSensitiveInfo(input);
113+
assertEquals(expectedOutput, actualOutput);
114+
}
115+
116+
@Test
117+
public void testMaskSensitiveInfo_noCredential() {
118+
String input = "{\"otherInfo\":\"someValue\"}";
119+
String expectedOutput = "{\"otherInfo\":\"someValue\"}";
120+
String actualOutput = maskSensitiveInfo(input);
121+
assertEquals(expectedOutput, actualOutput);
122+
}
123+
124+
@Test
125+
public void testMaskSensitiveInfo_emptyInput() {
126+
String input = "";
127+
String expectedOutput = "";
128+
String actualOutput = maskSensitiveInfo(input);
129+
assertEquals(expectedOutput, actualOutput);
104130
}
105131

106132
public void testSearchConnectors_beforeConnectorCreation() throws IOException {
107133
String searchEntity = "{\n" + " \"query\": {\n" + " \"match_all\": {}\n" + " },\n" + " \"size\": 1000\n" + "}";
108134
Response response = TestHelper
109135
.makeRequest(client(), "GET", "/_plugins/_ml/connectors/_search", null, TestHelper.toHttpEntity(searchEntity), null);
110136
Map responseMap = parseResponseToMap(response);
111-
assertEquals((Double) 0.0, (Double) ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
137+
assertEquals(0.0, ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
112138
}
113139

114140
public void testSearchConnectors_afterCreation() throws IOException {
@@ -125,7 +151,7 @@ public void testSearchRemoteModels_beforeCreation() throws IOException {
125151
Response response = TestHelper
126152
.makeRequest(client(), "GET", "/_plugins/_ml/models/_search", null, TestHelper.toHttpEntity(searchEntity), null);
127153
Map responseMap = parseResponseToMap(response);
128-
assertEquals((Double) 0.0, (Double) ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
154+
assertEquals(0.0, ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
129155
}
130156

131157
public void testSearchRemoteModels_afterCreation() throws IOException {
@@ -134,15 +160,15 @@ public void testSearchRemoteModels_afterCreation() throws IOException {
134160
Response response = TestHelper
135161
.makeRequest(client(), "GET", "/_plugins/_ml/models/_search", null, TestHelper.toHttpEntity(searchEntity), null);
136162
Map responseMap = parseResponseToMap(response);
137-
assertEquals((Double) 1.0, (Double) ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
163+
assertEquals(1.0, ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
138164
}
139165

140166
public void testSearchModelGroups_beforeCreation() throws IOException {
141167
String searchEntity = "{\n" + " \"query\": {\n" + " \"match_all\": {}\n" + " },\n" + " \"size\": 1000\n" + "}";
142168
Response response = TestHelper
143169
.makeRequest(client(), "GET", "/_plugins/_ml/model_groups/_search", null, TestHelper.toHttpEntity(searchEntity), null);
144170
Map responseMap = parseResponseToMap(response);
145-
assertEquals((Double) 0.0, (Double) ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
171+
assertEquals(0.0, ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
146172
}
147173

148174
public void testSearchModelGroups_afterCreation() throws IOException {
@@ -151,15 +177,15 @@ public void testSearchModelGroups_afterCreation() throws IOException {
151177
Response response = TestHelper
152178
.makeRequest(client(), "GET", "/_plugins/_ml/model_groups/_search", null, TestHelper.toHttpEntity(searchEntity), null);
153179
Map responseMap = parseResponseToMap(response);
154-
assertEquals((Double) 1.0, (Double) ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
180+
assertEquals(1.0, ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
155181
}
156182

157183
public void testSearchMLTasks_beforeCreation() throws IOException {
158184
String searchEntity = "{\n" + " \"query\": {\n" + " \"match_all\": {}\n" + " },\n" + " \"size\": 1000\n" + "}";
159185
Response response = TestHelper
160186
.makeRequest(client(), "GET", "/_plugins/_ml/tasks/_search", null, TestHelper.toHttpEntity(searchEntity), null);
161187
Map responseMap = parseResponseToMap(response);
162-
assertEquals((Double) 0.0, (Double) ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
188+
assertEquals(0.0, ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
163189
}
164190

165191
public void testSearchMLTasks_afterCreation() throws IOException {
@@ -169,7 +195,7 @@ public void testSearchMLTasks_afterCreation() throws IOException {
169195
Response response = TestHelper
170196
.makeRequest(client(), "GET", "/_plugins/_ml/tasks/_search", null, TestHelper.toHttpEntity(searchEntity), null);
171197
Map responseMap = parseResponseToMap(response);
172-
assertEquals((Double) 1.0, (Double) ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
198+
assertEquals(1.0, ((Map) ((Map) responseMap.get("hits")).get("total")).get("value"));
173199
}
174200

175201
public void testDeployRemoteModel() throws IOException, InterruptedException {
@@ -185,7 +211,7 @@ public void testDeployRemoteModel() throws IOException, InterruptedException {
185211
String modelId = (String) responseMap.get("model_id");
186212
response = deployRemoteModel(modelId);
187213
responseMap = parseResponseToMap(response);
188-
assertEquals("COMPLETED", (String) responseMap.get("status"));
214+
assertEquals("COMPLETED", responseMap.get("status"));
189215
taskId = (String) responseMap.get("task_id");
190216
waitForTask(taskId, MLTaskState.COMPLETED);
191217
}
@@ -838,7 +864,12 @@ public void testCohereClassifyModel() throws IOException, InterruptedException {
838864
}
839865

840866
public static Response createConnector(String input) throws IOException {
841-
return TestHelper.makeRequest(client(), "POST", "/_plugins/_ml/connectors/_create", null, TestHelper.toHttpEntity(input), null);
867+
try {
868+
return TestHelper.makeRequest(client(), "POST", "/_plugins/_ml/connectors/_create", null, TestHelper.toHttpEntity(input), null);
869+
} catch (ResponseException e) {
870+
String sanitizedMessage = maskSensitiveInfo(e.getMessage());// Log sanitized message
871+
throw new RuntimeException("Request failed: " + sanitizedMessage); // Re-throw sanitized exception
872+
}
842873
}
843874

844875
public static Response registerRemoteModel(String name, String connectorId) throws IOException {

0 commit comments

Comments
 (0)