18
18
import static org .opensearch .ml .action .models .DeleteModelTransportAction .OS_STATUS_EXCEPTION_MESSAGE ;
19
19
import static org .opensearch .ml .action .models .DeleteModelTransportAction .SEARCH_FAILURE_MSG ;
20
20
import static org .opensearch .ml .action .models .DeleteModelTransportAction .TIMEOUT_MSG ;
21
+ import static org .opensearch .ml .common .CommonValue .ML_AGENT_INDEX ;
21
22
import static org .opensearch .ml .common .CommonValue .ML_MODEL_INDEX ;
22
23
24
+ import java .io .ByteArrayOutputStream ;
23
25
import java .io .IOException ;
26
+ import java .io .ObjectOutputStream ;
27
+ import java .nio .ByteBuffer ;
28
+ import java .nio .charset .StandardCharsets ;
24
29
import java .util .ArrayList ;
25
30
import java .util .Arrays ;
31
+ import java .util .List ;
32
+ import java .util .Map ;
26
33
34
+ import org .apache .lucene .search .TotalHits ;
27
35
import org .junit .Before ;
28
36
import org .junit .Ignore ;
29
37
import org .junit .Rule ;
36
44
import org .opensearch .action .bulk .BulkItemResponse ;
37
45
import org .opensearch .action .delete .DeleteResponse ;
38
46
import org .opensearch .action .get .GetResponse ;
47
+ import org .opensearch .action .ingest .GetPipelineAction ;
48
+ import org .opensearch .action .ingest .GetPipelineResponse ;
49
+ import org .opensearch .action .search .GetSearchPipelineAction ;
50
+ import org .opensearch .action .search .GetSearchPipelineResponse ;
51
+ import org .opensearch .action .search .SearchRequest ;
52
+ import org .opensearch .action .search .SearchResponse ;
39
53
import org .opensearch .action .support .ActionFilters ;
40
54
import org .opensearch .client .Client ;
55
+ import org .opensearch .client .Response ;
41
56
import org .opensearch .cluster .service .ClusterService ;
42
57
import org .opensearch .common .settings .Settings ;
43
58
import org .opensearch .common .util .concurrent .ThreadContext ;
44
59
import org .opensearch .common .xcontent .XContentFactory ;
60
+ import org .opensearch .common .xcontent .XContentType ;
45
61
import org .opensearch .core .action .ActionListener ;
62
+ import org .opensearch .core .action .ActionResponse ;
63
+ import org .opensearch .core .common .bytes .BytesArray ;
46
64
import org .opensearch .core .common .bytes .BytesReference ;
65
+ import org .opensearch .core .xcontent .MediaType ;
66
+ import org .opensearch .core .xcontent .MediaTypeRegistry ;
47
67
import org .opensearch .core .xcontent .NamedXContentRegistry ;
48
68
import org .opensearch .core .xcontent .ToXContent ;
49
69
import org .opensearch .core .xcontent .XContentBuilder ;
58
78
import org .opensearch .ml .engine .tools .RelatedModelIdHelper ;
59
79
import org .opensearch .ml .helper .ModelAccessControlHelper ;
60
80
import org .opensearch .ml .model .MLModelManager ;
81
+ import org .opensearch .search .SearchHit ;
82
+ import org .opensearch .search .SearchHits ;
83
+ import org .opensearch .search .pipeline .PipelineConfiguration ;
61
84
import org .opensearch .test .OpenSearchTestCase ;
62
85
import org .opensearch .threadpool .ThreadPool ;
63
86
import org .opensearch .transport .TransportService ;
@@ -84,6 +107,9 @@ public class DeleteModelTransportActionTests extends OpenSearchTestCase {
84
107
@ Mock
85
108
BulkByScrollResponse bulkByScrollResponse ;
86
109
110
+ @ Mock
111
+ SearchResponse searchResponse ;
112
+
87
113
@ Mock
88
114
NamedXContentRegistry xContentRegistry ;
89
115
@@ -107,6 +133,20 @@ public class DeleteModelTransportActionTests extends OpenSearchTestCase {
107
133
@ Mock
108
134
private RelatedModelIdHelper relatedModelIdHelper ;
109
135
136
+ @ Mock
137
+ private GetSearchPipelineResponse getSearchPipelineResponse ;
138
+
139
+ @ Mock
140
+ private org .opensearch .search .pipeline .PipelineConfiguration searchPipelineConfiguration ;
141
+
142
+
143
+ @ Mock
144
+ GetPipelineResponse getIngestionPipelineResponse ;
145
+
146
+ private BulkByScrollResponse emptyBulkByScrollResponse ;
147
+
148
+ private Map <String , Object > configDataMap ;
149
+
110
150
@ Before
111
151
public void setup () throws IOException {
112
152
MockitoAnnotations .openMocks (this );
@@ -137,6 +177,7 @@ public void setup() throws IOException {
137
177
when (clusterService .getSettings ()).thenReturn (settings );
138
178
when (client .threadPool ()).thenReturn (threadPool );
139
179
when (threadPool .getThreadContext ()).thenReturn (threadContext );
180
+ prepare ();
140
181
}
141
182
142
183
public void testDeleteModel_Success () throws IOException {
@@ -146,13 +187,6 @@ public void testDeleteModel_Success() throws IOException {
146
187
return null ;
147
188
}).when (client ).delete (any (), any ());
148
189
149
- doAnswer (invocation -> {
150
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
151
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
152
- listener .onResponse (response );
153
- return null ;
154
- }).when (client ).execute (any (), any (), any ());
155
-
156
190
GetResponse getResponse = prepareMLModel (MLModelState .REGISTERED , null , false );
157
191
doAnswer (invocation -> {
158
192
ActionListener <GetResponse > actionListener = invocation .getArgument (1 );
@@ -164,20 +198,63 @@ public void testDeleteModel_Success() throws IOException {
164
198
verify (actionListener ).onResponse (deleteResponse );
165
199
}
166
200
201
+ public void testDeleteModel_BlockedBySearchPipeline () throws IOException {
202
+ //org.opensearch.search.pipeline.PipelineConfiguration pipelineConfiguration = new PipelineConfiguration();
203
+ when (searchPipelineConfiguration .getId ()).thenReturn ("1" );
204
+ when (searchPipelineConfiguration .getConfigAsMap ()).thenReturn (configDataMap );
205
+ when (getSearchPipelineResponse .pipelines ()).thenReturn (List .of (searchPipelineConfiguration ));
206
+ doAnswer (invocation -> {
207
+ ActionListener <GetSearchPipelineResponse > listener = invocation .getArgument (2 );
208
+ listener .onResponse (getSearchPipelineResponse );
209
+ return null ;
210
+ }).when (client ).execute (eq (GetSearchPipelineAction .INSTANCE ), any (), any ());
211
+
212
+ deleteModelTransportAction .doExecute (null , mlModelDeleteRequest , actionListener );
213
+ ArgumentCaptor <Exception > argumentCaptor = ArgumentCaptor .forClass (Exception .class );
214
+ verify (actionListener ).onFailure (argumentCaptor .capture ());
215
+ assertEquals ("1 search pipelines are still using this model, please delete or update the pipelines first: [1]" , argumentCaptor .getValue ().getMessage ());
216
+ }
217
+
218
+ public void testDeleteModel_BlockedByIngestPipeline () throws IOException {
219
+ org .opensearch .ingest .PipelineConfiguration ingestPipelineConfiguration = new org .opensearch .ingest .PipelineConfiguration (
220
+ "1" , new BytesArray ("{\" model_id\" : \" test_id\" }" .getBytes (StandardCharsets .UTF_8 )), MediaTypeRegistry .JSON
221
+ );
222
+ when (getIngestionPipelineResponse .pipelines ()).thenReturn (List .of (ingestPipelineConfiguration ));
223
+ doAnswer (invocation -> {
224
+ ActionListener <GetPipelineResponse > listener = invocation .getArgument (2 );
225
+ listener .onResponse (getIngestionPipelineResponse );
226
+ return null ;
227
+ }).when (client ).execute (eq (GetPipelineAction .INSTANCE ), any (), any ());
228
+
229
+ deleteModelTransportAction .doExecute (null , mlModelDeleteRequest , actionListener );
230
+ ArgumentCaptor <Exception > argumentCaptor = ArgumentCaptor .forClass (Exception .class );
231
+ verify (actionListener ).onFailure (argumentCaptor .capture ());
232
+ assertEquals ("1 ingest pipelines are still using this model, please delete or update the pipelines first: [1]" , argumentCaptor .getValue ().getMessage ());
233
+ }
234
+
235
+ public void testDeleteModel_BlockedByAgent () throws IOException {
236
+ SearchHit hit = new SearchHit (1 , "1" , null , null );
237
+ SearchHits searchHits = new SearchHits (new SearchHit [] { hit }, new TotalHits (1 , TotalHits .Relation .EQUAL_TO ), 1.0f );
238
+ when (searchResponse .getHits ()).thenReturn (searchHits );
239
+ doAnswer (invocation -> {
240
+ ActionListener <SearchResponse > listener = invocation .getArgument (1 );
241
+ listener .onResponse (searchResponse );
242
+ return null ;
243
+ }).when (client ).search (any (), any ());
244
+
245
+ deleteModelTransportAction .doExecute (null , mlModelDeleteRequest , actionListener );
246
+ ArgumentCaptor <Exception > argumentCaptor = ArgumentCaptor .forClass (Exception .class );
247
+ verify (actionListener ).onFailure (argumentCaptor .capture ());
248
+ assertEquals ("1 agents are still using this model, please delete or update the agents first: [1]" , argumentCaptor .getValue ().getMessage ());
249
+ }
250
+
167
251
public void testDeleteRemoteModel_Success () throws IOException {
168
252
doAnswer (invocation -> {
169
253
ActionListener <DeleteResponse > listener = invocation .getArgument (1 );
170
254
listener .onResponse (deleteResponse );
171
255
return null ;
172
256
}).when (client ).delete (any (), any ());
173
257
174
- doAnswer (invocation -> {
175
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
176
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
177
- listener .onResponse (response );
178
- return null ;
179
- }).when (client ).execute (any (), any (), any ());
180
-
181
258
GetResponse getResponse = prepareModelWithFunction (MLModelState .REGISTERED , null , false , FunctionName .REMOTE );
182
259
doAnswer (invocation -> {
183
260
ActionListener <GetResponse > actionListener = invocation .getArgument (1 );
@@ -200,12 +277,6 @@ public void testDeleteRemoteModel_deleteModelController_failed() throws IOExcept
200
277
return null ;
201
278
}).when (client ).delete (any (), any ());
202
279
203
- doAnswer (invocation -> {
204
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
205
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
206
- listener .onResponse (response );
207
- return null ;
208
- }).when (client ).execute (any (), any (), any ());
209
280
210
281
GetResponse getResponse = prepareModelWithFunction (MLModelState .REGISTERED , null , false , FunctionName .REMOTE );
211
282
doAnswer (invocation -> {
@@ -231,13 +302,6 @@ public void testDeleteLocalModel_deleteModelController_failed() throws IOExcepti
231
302
return null ;
232
303
}).when (client ).delete (any (), any ());
233
304
234
- doAnswer (invocation -> {
235
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
236
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
237
- listener .onResponse (response );
238
- return null ;
239
- }).when (client ).execute (any (), any (), any ());
240
-
241
305
GetResponse getResponse = prepareModelWithFunction (MLModelState .REGISTERED , null , false , FunctionName .TEXT_EMBEDDING );
242
306
doAnswer (invocation -> {
243
307
ActionListener <GetResponse > actionListener = invocation .getArgument (1 );
@@ -284,13 +348,6 @@ public void testDeleteHiddenModel_Success() throws IOException {
284
348
return null ;
285
349
}).when (client ).delete (any (), any ());
286
350
287
- doAnswer (invocation -> {
288
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
289
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
290
- listener .onResponse (response );
291
- return null ;
292
- }).when (client ).execute (any (), any (), any ());
293
-
294
351
GetResponse getResponse = prepareMLModel (MLModelState .REGISTERED , null , true );
295
352
doAnswer (invocation -> {
296
353
ActionListener <GetResponse > actionListener = invocation .getArgument (1 );
@@ -310,13 +367,6 @@ public void testDeleteHiddenModel_NoSuperAdminPermission() throws IOException {
310
367
return null ;
311
368
}).when (client ).delete (any (), any ());
312
369
313
- doAnswer (invocation -> {
314
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
315
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
316
- listener .onResponse (response );
317
- return null ;
318
- }).when (client ).execute (any (), any (), any ());
319
-
320
370
GetResponse getResponse = prepareMLModel (MLModelState .REGISTERED , null , true );
321
371
doAnswer (invocation -> {
322
372
ActionListener <GetResponse > actionListener = invocation .getArgument (1 );
@@ -338,13 +388,6 @@ public void testDeleteModel_Success_AlgorithmNotNull() throws IOException {
338
388
return null ;
339
389
}).when (client ).delete (any (), any ());
340
390
341
- doAnswer (invocation -> {
342
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
343
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
344
- listener .onResponse (response );
345
- return null ;
346
- }).when (client ).execute (any (), any (), any ());
347
-
348
391
GetResponse getResponse = prepareMLModel (MLModelState .REGISTERED , null , false );
349
392
doAnswer (invocation -> {
350
393
ActionListener <GetResponse > actionListener = invocation .getArgument (1 );
@@ -417,12 +460,6 @@ public void testDeleteModel_deleteModelController_ResourceNotFoundException() th
417
460
return null ;
418
461
}).when (client ).delete (any (), any ());
419
462
420
- doAnswer (invocation -> {
421
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
422
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
423
- listener .onResponse (response );
424
- return null ;
425
- }).when (client ).execute (any (), any (), any ());
426
463
427
464
GetResponse getResponse = prepareMLModel (MLModelState .REGISTERED , null , false );
428
465
doAnswer (invocation -> {
@@ -467,12 +504,6 @@ public void testDeleteRemoteModel_modelNotFound_ResourceNotFoundException() thro
467
504
return null ;
468
505
}).when (client ).delete (any (), any ());
469
506
470
- doAnswer (invocation -> {
471
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
472
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
473
- listener .onResponse (response );
474
- return null ;
475
- }).when (client ).execute (any (), any (), any ());
476
507
477
508
GetResponse getResponse = prepareModelWithFunction (MLModelState .REGISTERED , null , false , FunctionName .REMOTE );
478
509
doAnswer (invocation -> {
@@ -498,12 +529,6 @@ public void testDeleteRemoteModel_modelNotFound_RuntimeException() throws IOExce
498
529
return null ;
499
530
}).when (client ).delete (any (), any ());
500
531
501
- doAnswer (invocation -> {
502
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
503
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
504
- listener .onResponse (response );
505
- return null ;
506
- }).when (client ).execute (any (), any (), any ());
507
532
508
533
GetResponse getResponse = prepareModelWithFunction (MLModelState .REGISTERED , null , false , FunctionName .REMOTE );
509
534
doAnswer (invocation -> {
@@ -531,12 +556,6 @@ public void testModelNotFound_modelChunks_modelController_delete_success() throw
531
556
return null ;
532
557
}).when (client ).delete (any (), any ());
533
558
534
- doAnswer (invocation -> {
535
- ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
536
- BulkByScrollResponse response = new BulkByScrollResponse (new ArrayList <>(), null );
537
- listener .onResponse (response );
538
- return null ;
539
- }).when (client ).execute (any (), any (), any ());
540
559
deleteModelTransportAction .doExecute (null , mlModelDeleteRequest , actionListener );
541
560
ArgumentCaptor <OpenSearchStatusException > argumentCaptor = ArgumentCaptor .forClass (OpenSearchStatusException .class );
542
561
verify (actionListener ).onFailure (argumentCaptor .capture ());
@@ -661,4 +680,40 @@ private GetResponse buildResponse(MLModel mlModel) throws IOException {
661
680
GetResponse getResponse = new GetResponse (getResult );
662
681
return getResponse ;
663
682
}
683
+
684
+ private void prepare () {
685
+ emptyBulkByScrollResponse = new BulkByScrollResponse (new ArrayList <>(), null );
686
+ SearchHits hits = new SearchHits (new SearchHit [] { }, new TotalHits (0 , TotalHits .Relation .EQUAL_TO ), 0.0f );
687
+ when (searchResponse .getHits ()).thenReturn (hits );
688
+ when (getIngestionPipelineResponse .pipelines ()).thenReturn (List .of ());
689
+
690
+ doAnswer (invocation -> {
691
+ ActionListener <SearchResponse > listener = invocation .getArgument (1 );
692
+ listener .onResponse (searchResponse );
693
+ return null ;
694
+ }).when (client ).search (any (), any ());
695
+
696
+
697
+ doAnswer (invocation -> {
698
+ ActionListener <BulkByScrollResponse > listener = invocation .getArgument (2 );
699
+ listener .onResponse (emptyBulkByScrollResponse );
700
+ return null ;
701
+ }).when (client ).execute (eq (DeleteByQueryAction .INSTANCE ), any (), any ());
702
+
703
+ doAnswer (invocation -> {
704
+ ActionListener <GetPipelineResponse > listener = invocation .getArgument (2 );
705
+ listener .onResponse (getIngestionPipelineResponse );
706
+ return null ;
707
+ }).when (client ).execute (eq (GetPipelineAction .INSTANCE ), any (), any ());
708
+
709
+ when (getSearchPipelineResponse .pipelines ()).thenReturn (List .of ());
710
+ doAnswer (invocation -> {
711
+ ActionListener <GetSearchPipelineResponse > listener = invocation .getArgument (2 );
712
+ listener .onResponse (getSearchPipelineResponse );
713
+ return null ;
714
+ }).when (client ).execute (eq (GetSearchPipelineAction .INSTANCE ), any (), any ());
715
+ configDataMap = Map .of ("model_id" , "test_id" , "list_model_id" , List .of ("test_list_id" ),
716
+ "test_map_id" , Map .of ("test_key" , "test_map_id" ));
717
+ doAnswer (invocation -> new SearchRequest ()).when (relatedModelIdHelper ).constructQueryRequest (any ());
718
+ }
664
719
}
0 commit comments