Skip to content

Commit d9e0925

Browse files
committed
add test for cluster setting
Signed-off-by: xinyual <xinyual@amazon.com>
1 parent 8ef451e commit d9e0925

File tree

1 file changed

+76
-0
lines changed

1 file changed

+76
-0
lines changed

plugin/src/test/java/org/opensearch/ml/action/models/DeleteModelTransportActionTests.java

+76
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,20 @@ public void testDeleteModel_Success() throws IOException {
198198
}
199199

200200
public void testDeleteModel_BlockedBySearchPipelineAndIngestionPipeline() throws IOException {
201+
doAnswer(invocation -> {
202+
ActionListener<DeleteResponse> listener = invocation.getArgument(1);
203+
listener.onResponse(deleteResponse);
204+
return null;
205+
}).when(client).delete(any(), any());
206+
207+
GetResponse getResponse = prepareMLModel(MLModelState.REGISTERED, null, false);
208+
doAnswer(invocation -> {
209+
ActionListener<GetResponse> actionListener = invocation.getArgument(1);
210+
actionListener.onResponse(getResponse);
211+
return null;
212+
}).when(client).get(any(), any());
213+
214+
201215
when(getSearchPipelineResponse.toString()).thenReturn(
202216
StringUtils.toJson(Map.of("search_1", configDataMap))
203217
);
@@ -290,6 +304,68 @@ public void testDeleteModel_BlockedBySearchPipelineListModelId() throws IOExcept
290304
);
291305
}
292306

307+
public void testDeleteModel_UseSettingToSkipBlockedByIngestPipeline() throws IOException {
308+
Settings settings = Settings.builder().put(ML_COMMONS_SAFE_DELETE_MODEL.getKey(), false).build();
309+
threadContext = new ThreadContext(settings);
310+
ClusterSettings clusterSettings = clusterSetting(settings, ML_COMMONS_SAFE_DELETE_MODEL);
311+
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
312+
when(clusterService.getSettings()).thenReturn(settings);
313+
deleteModelTransportAction = spy(
314+
new DeleteModelTransportAction(
315+
transportService,
316+
actionFilters,
317+
client,
318+
settings,
319+
xContentRegistry,
320+
clusterService,
321+
modelAccessControlHelper,
322+
agentModelsSearcher
323+
)
324+
);
325+
326+
threadContext = new ThreadContext(settings);
327+
when(clusterService.getSettings()).thenReturn(settings);
328+
when(client.threadPool()).thenReturn(threadPool);
329+
when(threadPool.getThreadContext()).thenReturn(threadContext);
330+
331+
doAnswer(invocation -> {
332+
ActionListener<Boolean> listener = invocation.getArgument(3);
333+
listener.onResponse(true);
334+
return null;
335+
}).when(modelAccessControlHelper).validateModelGroupAccess(any(), any(), any(), any());
336+
337+
338+
339+
Map<String, Object> ingestPipelineConfig1 = Map.of("model_id", "test_id");
340+
Map<String, Object> ingestPipelineConfig2 = Map.of("nothing", "test_id");
341+
when(getIngestionPipelineResponse.toString())
342+
.thenReturn(StringUtils.toJson(Map.of("ingest_1", ingestPipelineConfig1, "ingest_2", ingestPipelineConfig2)));
343+
// when(getIngestionPipelineResponse.pipelines())
344+
// .thenReturn(List.of(ingestPipelineConfiguration, independentIngestPipelineConfiguration));
345+
doAnswer(invocation -> {
346+
ActionListener<GetPipelineResponse> listener = invocation.getArgument(2);
347+
listener.onResponse(getIngestionPipelineResponse);
348+
return null;
349+
}).when(client).execute(eq(GetPipelineAction.INSTANCE), any(), any());
350+
351+
doAnswer(invocation -> {
352+
ActionListener<DeleteResponse> listener = invocation.getArgument(1);
353+
listener.onResponse(deleteResponse);
354+
return null;
355+
}).when(client).delete(any(), any());
356+
357+
GetResponse getResponse = prepareMLModel(MLModelState.REGISTERED, null, false);
358+
doAnswer(invocation -> {
359+
ActionListener<GetResponse> actionListener = invocation.getArgument(1);
360+
actionListener.onResponse(getResponse);
361+
return null;
362+
}).when(client).get(any(), any());
363+
364+
deleteModelTransportAction.doExecute(null, mlModelDeleteRequest, actionListener);
365+
verify(actionListener).onResponse(deleteResponse);
366+
367+
}
368+
293369
public void testDeleteModel_BlockedByIngestPipeline() throws IOException {
294370
Map<String, Object> ingestPipelineConfig1 = Map.of("model_id", "test_id");
295371
Map<String, Object> ingestPipelineConfig2 = Map.of("nothing", "test_id");

0 commit comments

Comments
 (0)