66
77package org .elasticsearch .xpack .core .ml ;
88
9+ import org .elasticsearch .Version ;
10+ import org .elasticsearch .cluster .node .DiscoveryNode ;
11+ import org .elasticsearch .cluster .node .DiscoveryNodes ;
12+ import org .elasticsearch .common .transport .TransportAddress ;
913import org .elasticsearch .persistent .PersistentTasksCustomMetaData ;
1014import org .elasticsearch .test .ESTestCase ;
1115import org .elasticsearch .xpack .core .ml .action .OpenJobAction ;
1418import org .elasticsearch .xpack .core .ml .job .config .JobState ;
1519import org .elasticsearch .xpack .core .ml .job .config .JobTaskState ;
1620
21+ import java .net .InetAddress ;
22+
1723import static org .hamcrest .Matchers .containsInAnyOrder ;
1824import static org .hamcrest .Matchers .empty ;
1925
2026public class MlTasksTests extends ESTestCase {
2127 public void testGetJobState () {
22- PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
28+ PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
2329 // A missing task is a closed job
2430 assertEquals (JobState .CLOSED , MlTasks .getJobState ("foo" , tasksBuilder .build ()));
2531 // A task with no status is opening
@@ -52,7 +58,7 @@ public void testGetDatefeedState() {
5258 public void testGetJobTask () {
5359 assertNull (MlTasks .getJobTask ("foo" , null ));
5460
55- PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
61+ PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
5662 tasksBuilder .addTask (MlTasks .jobTaskId ("foo" ), MlTasks .JOB_TASK_NAME , new OpenJobAction .JobParams ("foo" ),
5763 new PersistentTasksCustomMetaData .Assignment ("bar" , "test assignment" ));
5864
@@ -73,7 +79,7 @@ public void testGetDatafeedTask() {
7379 }
7480
7581 public void testOpenJobIds () {
76- PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
82+ PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
7783 assertThat (MlTasks .openJobIds (tasksBuilder .build ()), empty ());
7884
7985 tasksBuilder .addTask (MlTasks .jobTaskId ("foo-1" ), MlTasks .JOB_TASK_NAME , new OpenJobAction .JobParams ("foo-1" ),
@@ -92,7 +98,7 @@ public void testOpenJobIds_GivenNull() {
9298 }
9399
94100 public void testStartedDatafeedIds () {
95- PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
101+ PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
96102 assertThat (MlTasks .openJobIds (tasksBuilder .build ()), empty ());
97103
98104 tasksBuilder .addTask (MlTasks .jobTaskId ("job-1" ), MlTasks .JOB_TASK_NAME , new OpenJobAction .JobParams ("foo-1" ),
@@ -111,16 +117,48 @@ public void testStartedDatafeedIds_GivenNull() {
111117 assertThat (MlTasks .startedDatafeedIds (null ), empty ());
112118 }
113119
114- public void testTaskExistsForJob () {
115- PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
116- assertFalse (MlTasks .taskExistsForJob ("job-1" , tasksBuilder .build ()));
117-
118- tasksBuilder .addTask (MlTasks .jobTaskId ("foo" ), MlTasks .JOB_TASK_NAME , new OpenJobAction .JobParams ("foo" ),
119- new PersistentTasksCustomMetaData .Assignment ("node-1" , "test assignment" ));
120- tasksBuilder .addTask (MlTasks .jobTaskId ("bar" ), MlTasks .JOB_TASK_NAME , new OpenJobAction .JobParams ("bar" ),
120+ public void testUnallocatedJobIds () {
121+ PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
122+ tasksBuilder .addTask (MlTasks .jobTaskId ("job_with_assignment" ), MlTasks .JOB_TASK_NAME ,
123+ new OpenJobAction .JobParams ("job_with_assignment" ),
121124 new PersistentTasksCustomMetaData .Assignment ("node-1" , "test assignment" ));
125+ tasksBuilder .addTask (MlTasks .jobTaskId ("job_without_assignment" ), MlTasks .JOB_TASK_NAME ,
126+ new OpenJobAction .JobParams ("job_without_assignment" ),
127+ new PersistentTasksCustomMetaData .Assignment (null , "test assignment" ));
128+ tasksBuilder .addTask (MlTasks .jobTaskId ("job_without_node" ), MlTasks .JOB_TASK_NAME ,
129+ new OpenJobAction .JobParams ("job_without_node" ),
130+ new PersistentTasksCustomMetaData .Assignment ("dead-node" , "expired node" ));
131+
132+ DiscoveryNodes nodes = DiscoveryNodes .builder ()
133+ .add (new DiscoveryNode ("node-1" , new TransportAddress (InetAddress .getLoopbackAddress (), 9300 ), Version .CURRENT ))
134+ .localNodeId ("node-1" )
135+ .masterNodeId ("node-1" )
136+ .build ();
137+
138+ assertThat (MlTasks .unallocatedJobIds (tasksBuilder .build (), nodes ),
139+ containsInAnyOrder ("job_without_assignment" , "job_without_node" ));
140+ }
122141
123- assertFalse (MlTasks .taskExistsForJob ("job-1" , tasksBuilder .build ()));
124- assertTrue (MlTasks .taskExistsForJob ("foo" , tasksBuilder .build ()));
142+ public void testUnallocatedDatafeedIds () {
143+ PersistentTasksCustomMetaData .Builder tasksBuilder = PersistentTasksCustomMetaData .builder ();
144+ tasksBuilder .addTask (MlTasks .datafeedTaskId ("datafeed_with_assignment" ), MlTasks .DATAFEED_TASK_NAME ,
145+ new StartDatafeedAction .DatafeedParams ("datafeed_with_assignment" , 0L ),
146+ new PersistentTasksCustomMetaData .Assignment ("node-1" , "test assignment" ));
147+ tasksBuilder .addTask (MlTasks .datafeedTaskId ("datafeed_without_assignment" ), MlTasks .DATAFEED_TASK_NAME ,
148+ new StartDatafeedAction .DatafeedParams ("datafeed_without_assignment" , 0L ),
149+ new PersistentTasksCustomMetaData .Assignment (null , "test assignment" ));
150+ tasksBuilder .addTask (MlTasks .datafeedTaskId ("datafeed_without_node" ), MlTasks .DATAFEED_TASK_NAME ,
151+ new StartDatafeedAction .DatafeedParams ("datafeed_without_node" , 0L ),
152+ new PersistentTasksCustomMetaData .Assignment ("dead_node" , "expired node" ));
153+
154+
155+ DiscoveryNodes nodes = DiscoveryNodes .builder ()
156+ .add (new DiscoveryNode ("node-1" , new TransportAddress (InetAddress .getLoopbackAddress (), 9300 ), Version .CURRENT ))
157+ .localNodeId ("node-1" )
158+ .masterNodeId ("node-1" )
159+ .build ();
160+
161+ assertThat (MlTasks .unallocatedDatafeedIds (tasksBuilder .build (), nodes ),
162+ containsInAnyOrder ("datafeed_without_assignment" , "datafeed_without_node" ));
125163 }
126164}
0 commit comments