1717import org .junit .Before ;
1818import org .junit .Ignore ;
1919import org .junit .Rule ;
20+ import org .junit .Test ;
2021import org .junit .rules .ExpectedException ;
2122import org .opensearch .client .Response ;
23+ import org .opensearch .client .ResponseException ;
2224import org .opensearch .ml .common .MLTaskState ;
2325import 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