|
14 | 14 |
|
15 | 15 | import java.io.IOException; |
16 | 16 | import java.nio.file.Path; |
| 17 | +import java.time.Instant; |
17 | 18 | import java.util.Arrays; |
18 | 19 | import java.util.List; |
19 | 20 | import java.util.Map; |
@@ -202,6 +203,117 @@ public void setup() throws IOException { |
202 | 203 | }).when(connectorAccessControlHelper).getConnector(any(Client.class), any(String.class), isA(ActionListener.class)); |
203 | 204 | } |
204 | 205 |
|
| 206 | + @Test |
| 207 | + public void testUpdateConnectorDoesNotUpdateHttpConnectorTimeFields() { |
| 208 | + HttpConnector connector = HttpConnector |
| 209 | + .builder() |
| 210 | + .name("test") |
| 211 | + .protocol("http") |
| 212 | + .version("1") |
| 213 | + .credential(Map.of("api_key", "credential_value")) |
| 214 | + .parameters(Map.of("param1", "value1")) |
| 215 | + .actions( |
| 216 | + Arrays |
| 217 | + .asList( |
| 218 | + ConnectorAction |
| 219 | + .builder() |
| 220 | + .actionType(ConnectorAction.ActionType.PREDICT) |
| 221 | + .method("POST") |
| 222 | + .url("https://api.openai.com/v1/chat/completions") |
| 223 | + .headers(Map.of("Authorization", "Bearer ${credential.api_key}")) |
| 224 | + .requestBody("{ \"model\": \"${parameters.model}\", \"messages\": ${parameters.messages} }") |
| 225 | + .build() |
| 226 | + ) |
| 227 | + ) |
| 228 | + .build(); |
| 229 | + |
| 230 | + assertNull(connector.getCreatedTime()); |
| 231 | + assertNull(connector.getLastUpdateTime()); |
| 232 | + |
| 233 | + doReturn(true).when(connectorAccessControlHelper).validateConnectorAccess(any(Client.class), any(Connector.class)); |
| 234 | + |
| 235 | + doAnswer(invocation -> { |
| 236 | + ActionListener<Connector> listener = invocation.getArgument(2); |
| 237 | + listener.onResponse(connector); |
| 238 | + return null; |
| 239 | + }).when(connectorAccessControlHelper).getConnector(any(Client.class), any(String.class), isA(ActionListener.class)); |
| 240 | + |
| 241 | + doAnswer(invocation -> { |
| 242 | + ActionListener<SearchResponse> actionListener = invocation.getArgument(1); |
| 243 | + actionListener.onResponse(searchResponse); |
| 244 | + return null; |
| 245 | + }).when(client).search(any(SearchRequest.class), isA(ActionListener.class)); |
| 246 | + |
| 247 | + doAnswer(invocation -> { |
| 248 | + ActionListener<UpdateResponse> listener = invocation.getArgument(1); |
| 249 | + listener.onResponse(updateResponse); |
| 250 | + return null; |
| 251 | + }).when(client).update(any(UpdateRequest.class), isA(ActionListener.class)); |
| 252 | + |
| 253 | + updateConnectorTransportAction.doExecute(task, updateRequest, actionListener); |
| 254 | + |
| 255 | + assertNull(connector.getCreatedTime()); |
| 256 | + assertNotNull(connector.getLastUpdateTime()); |
| 257 | + } |
| 258 | + |
| 259 | + @Test |
| 260 | + public void testUpdateConnectorUpdatesHttpConnectorTimeFields() { |
| 261 | + HttpConnector connector = HttpConnector |
| 262 | + .builder() |
| 263 | + .name("test") |
| 264 | + .protocol("http") |
| 265 | + .version("1") |
| 266 | + .credential(Map.of("api_key", "credential_value")) |
| 267 | + .parameters(Map.of("param1", "value1")) |
| 268 | + .actions( |
| 269 | + Arrays |
| 270 | + .asList( |
| 271 | + ConnectorAction |
| 272 | + .builder() |
| 273 | + .actionType(ConnectorAction.ActionType.PREDICT) |
| 274 | + .method("POST") |
| 275 | + .url("https://api.openai.com/v1/chat/completions") |
| 276 | + .headers(Map.of("Authorization", "Bearer ${credential.api_key}")) |
| 277 | + .requestBody("{ \"model\": \"${parameters.model}\", \"messages\": ${parameters.messages} }") |
| 278 | + .build() |
| 279 | + ) |
| 280 | + ) |
| 281 | + .build(); |
| 282 | + |
| 283 | + Instant testInitialTime = Instant.now(); |
| 284 | + connector.setCreatedTime(testInitialTime); |
| 285 | + connector.setLastUpdateTime(testInitialTime); |
| 286 | + |
| 287 | + assert (connector.getCreatedTime().toEpochMilli() == connector.getLastUpdateTime().toEpochMilli()); |
| 288 | + |
| 289 | + doReturn(true).when(connectorAccessControlHelper).validateConnectorAccess(any(Client.class), any(Connector.class)); |
| 290 | + |
| 291 | + doAnswer(invocation -> { |
| 292 | + ActionListener<Connector> listener = invocation.getArgument(2); |
| 293 | + listener.onResponse(connector); |
| 294 | + return null; |
| 295 | + }).when(connectorAccessControlHelper).getConnector(any(Client.class), any(String.class), isA(ActionListener.class)); |
| 296 | + |
| 297 | + doAnswer(invocation -> { |
| 298 | + ActionListener<SearchResponse> actionListener = invocation.getArgument(1); |
| 299 | + actionListener.onResponse(searchResponse); |
| 300 | + return null; |
| 301 | + }).when(client).search(any(SearchRequest.class), isA(ActionListener.class)); |
| 302 | + |
| 303 | + doAnswer(invocation -> { |
| 304 | + ActionListener<UpdateResponse> listener = invocation.getArgument(1); |
| 305 | + listener.onResponse(updateResponse); |
| 306 | + return null; |
| 307 | + }).when(client).update(any(UpdateRequest.class), isA(ActionListener.class)); |
| 308 | + |
| 309 | + updateConnectorTransportAction.doExecute(task, updateRequest, actionListener); |
| 310 | + |
| 311 | + assertTrue( |
| 312 | + "Last update time must be bigger than the creation time", |
| 313 | + connector.getLastUpdateTime().toEpochMilli() >= connector.getCreatedTime().toEpochMilli() |
| 314 | + ); |
| 315 | + } |
| 316 | + |
205 | 317 | @Test |
206 | 318 | public void testExecuteConnectorAccessControlSuccess() { |
207 | 319 | doReturn(true).when(connectorAccessControlHelper).validateConnectorAccess(any(Client.class), any(Connector.class)); |
|
0 commit comments