2626import org .elasticsearch .test .ESIntegTestCase ;
2727import org .elasticsearch .test .discovery .TestZenDiscovery ;
2828import org .elasticsearch .xpack .core .XPackSettings ;
29+ import org .elasticsearch .xpack .core .ml .action .GetDatafeedsAction ;
30+ import org .elasticsearch .xpack .core .ml .action .GetJobsAction ;
31+ import org .elasticsearch .xpack .core .ml .action .util .QueryPage ;
32+ import org .elasticsearch .xpack .core .ml .client .MachineLearningClient ;
2933import org .elasticsearch .xpack .ml .LocalStateMachineLearning ;
3034import org .elasticsearch .xpack .ml .MachineLearning ;
3135import org .elasticsearch .xpack .core .ml .MachineLearningField ;
@@ -270,7 +274,9 @@ public static GetDatafeedsStatsAction.Response.DatafeedStats getDatafeedStats(St
270274 }
271275
272276 public static void deleteAllDatafeeds (Logger logger , Client client ) throws Exception {
273- MlMetadata mlMetadata = MlMetadata .getMlMetadata (client .admin ().cluster ().prepareState ().get ().getState ());
277+ final MachineLearningClient mlClient = new MachineLearningClient (client );
278+ final QueryPage <DatafeedConfig > datafeeds =
279+ mlClient .getDatafeeds (new GetDatafeedsAction .Request (GetDatafeedsAction .ALL )).actionGet ().getResponse ();
274280 try {
275281 logger .info ("Closing all datafeeds (using _all)" );
276282 StopDatafeedAction .Response stopResponse = client
@@ -291,25 +297,25 @@ public static void deleteAllDatafeeds(Logger logger, Client client) throws Excep
291297 "Had to resort to force-stopping datafeed, something went wrong?" , e1 );
292298 }
293299
294- for (DatafeedConfig datafeed : mlMetadata .getDatafeeds ().values ()) {
295- String datafeedId = datafeed .getId ();
300+ for (final DatafeedConfig datafeed : datafeeds .results ()) {
296301 assertBusy (() -> {
297302 try {
298- GetDatafeedsStatsAction .Request request = new GetDatafeedsStatsAction .Request (datafeedId );
303+ GetDatafeedsStatsAction .Request request = new GetDatafeedsStatsAction .Request (datafeed . getId () );
299304 GetDatafeedsStatsAction .Response r = client .execute (GetDatafeedsStatsAction .INSTANCE , request ).get ();
300305 assertThat (r .getResponse ().results ().get (0 ).getDatafeedState (), equalTo (DatafeedState .STOPPED ));
301306 } catch (InterruptedException | ExecutionException e ) {
302307 throw new RuntimeException (e );
303308 }
304309 });
305310 DeleteDatafeedAction .Response deleteResponse =
306- client .execute (DeleteDatafeedAction .INSTANCE , new DeleteDatafeedAction .Request (datafeedId )).get ();
311+ client .execute (DeleteDatafeedAction .INSTANCE , new DeleteDatafeedAction .Request (datafeed . getId () )).get ();
307312 assertTrue (deleteResponse .isAcknowledged ());
308313 }
309314 }
310315
311316 public static void deleteAllJobs (Logger logger , Client client ) throws Exception {
312- MlMetadata mlMetadata = MlMetadata .getMlMetadata (client .admin ().cluster ().prepareState ().get ().getState ());
317+ final MachineLearningClient mlClient = new MachineLearningClient (client );
318+ final QueryPage <Job > jobs = mlClient .getJobs (new GetJobsAction .Request (MetaData .ALL )).actionGet ().getResponse ();
313319
314320 try {
315321 CloseJobAction .Request closeRequest = new CloseJobAction .Request (MetaData .ALL );
@@ -333,15 +339,14 @@ public static void deleteAllJobs(Logger logger, Client client) throws Exception
333339 e1 );
334340 }
335341
336- for (Map .Entry <String , Job > entry : mlMetadata .getJobs ().entrySet ()) {
337- String jobId = entry .getKey ();
342+ for (final Job job : jobs .results ()) {
338343 assertBusy (() -> {
339344 GetJobsStatsAction .Response statsResponse =
340- client ().execute (GetJobsStatsAction .INSTANCE , new GetJobsStatsAction .Request (jobId )).actionGet ();
345+ client ().execute (GetJobsStatsAction .INSTANCE , new GetJobsStatsAction .Request (job . getId () )).actionGet ();
341346 assertEquals (JobState .CLOSED , statsResponse .getResponse ().results ().get (0 ).getState ());
342347 });
343348 DeleteJobAction .Response response =
344- client .execute (DeleteJobAction .INSTANCE , new DeleteJobAction .Request (jobId )).get ();
349+ client .execute (DeleteJobAction .INSTANCE , new DeleteJobAction .Request (job . getId () )).get ();
345350 assertTrue (response .isAcknowledged ());
346351 }
347352 }
0 commit comments