Unit tests commented and fixed

Author: valeriy42

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/process/autodetect/AutodetectProcessManager.java

@@ -1004,13 +1004,19 @@ public Optional<Duration> jobOpenTime(JobTask jobTask) {
 
     void setJobState(JobTask jobTask, JobState state, String reason) {
         JobTaskState jobTaskState = new JobTaskState(state, jobTask.getAllocationId(), reason, Instant.now());
-        jobTask.updatePersistentTaskState(
+        // retry state update to ensure that cluster state stays consistent
+        new UpdateStateRetryableAction(
+            logger,
+            threadPool,
+            TimeValue.timeValueMillis(UpdateStateRetryableAction.MIN_RETRY_SLEEP_MILLIS),
+            TimeValue.MAX_VALUE,
+            jobTask,
             jobTaskState,
             ActionListener.wrap(
                 persistentTask -> logger.info("Successfully set job state to [{}] for job [{}]", state, jobTask.getJobId()),
                 e -> logSetJobStateFailure(state, jobTask.getJobId(), e)
             )
-        );
+        ).run();
     }
 
     private static void logSetJobStateFailure(JobState state, String jobId, Exception e) {
@@ -1023,7 +1029,7 @@ private static void logSetJobStateFailure(JobState state, String jobId, Exceptio
 
     void setJobState(JobTask jobTask, JobState state, String reason, CheckedConsumer<Exception, IOException> handler) {
         JobTaskState jobTaskState = new JobTaskState(state, jobTask.getAllocationId(), reason, Instant.now());
-        // retry with a small initial backoff of 10ms
+        // retry state update to ensure that cluster state stays consistent
         new UpdateStateRetryableAction(
             logger,
             threadPool,

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/job/process/autodetect/AutodetectProcessManagerTests.java

@@ -869,27 +869,28 @@ public void testSetJobState_withoutHandler_invokesPersistentTaskUpdate() {
         // call the no-handler overload
         manager.setJobState(jobTask, JobState.CLOSING, "closing-reason");
 
-        // verify we created the correct JobTaskState and passed some listener
+        // verify we called updatePersistentTaskState with the expected state
         @SuppressWarnings("unchecked")
-        ArgumentCaptor<JobTaskState> stateCaptor =
-            ArgumentCaptor.forClass(JobTaskState.class);
+        ArgumentCaptor<JobTaskState> stateCaptor = ArgumentCaptor.forClass(JobTaskState.class);
         verify(jobTask).updatePersistentTaskState(stateCaptor.capture(), any());
         JobTaskState captured = stateCaptor.getValue();
         assertEquals(JobState.CLOSING, captured.getState());
         assertEquals(123L, captured.getAllocationId());
         assertEquals("closing-reason", captured.getReason());
-//        assertNotNull(captured.getTimestamp());
     }
 
     public void testSetJobState_withHandler_onResponse_triggersHandlerNull() throws IOException {
+        // This test verifies the “happy‐path” of the retryable overload—i.e. what happens when the very first call
+        // to updatePersistentTaskState succeeds. On a successful state update it must invoke handler.accept(null)
+        // (because there was no error).
         AutodetectProcessManager manager = createSpyManager();
         JobTask jobTask = mock(JobTask.class);
 
         // stub updatePersistentTaskState to call onResponse
         doAnswer(invocation -> {
             @SuppressWarnings("unchecked")
-            ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>> listener =
-                (ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>>) invocation.getArguments()[1];
+            ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>> listener = (ActionListener<
+                PersistentTasksCustomMetadata.PersistentTask<?>>) invocation.getArguments()[1];
             listener.onResponse(null);
             return null;
         }).when(jobTask).updatePersistentTaskState(any(), any());
@@ -905,11 +906,16 @@ public void testSetJobState_withHandler_onResponse_triggersHandlerNull() throws
     }
 
     public void testSetJobState_withHandler_onFailure_triggersHandlerException() throws IOException {
+        // Verifies that when updatePersistentTaskState reports a failure, the handler receives that exception
+        when(threadPool.schedule(any(Runnable.class), any(TimeValue.class), any(Executor.class)))
+            .thenAnswer(invocation -> {
+                Runnable r = invocation.getArgument(0);
+                r.run();
+                return mock(ThreadPool.Cancellable.class);
+            });
         AutodetectProcessManager manager = createSpyManager();
         JobTask jobTask = mock(JobTask.class);
-        Exception boom = new RuntimeException("boom");
-
-        // stub updatePersistentTaskState to call onFailure
+        ResourceNotFoundException boom = new ResourceNotFoundException("boom");
         doAnswer(invocation -> {
             @SuppressWarnings("unchecked")
             ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>> listener =
@@ -929,6 +935,9 @@ public void testSetJobState_withHandler_onFailure_triggersHandlerException() thr
     }
 
     public void testSetJobState_withHandler_retriesUntilSuccess() throws IOException {
+        // Verifies that transient failures are retried until eventual success, and the handler receives null on success
+
+        // ensure that all retries are executed on the same thread for determinism
         when(threadPool.schedule(any(Runnable.class), any(TimeValue.class), any(Executor.class))).thenAnswer(invocation -> {
             Runnable r = invocation.getArgument(0);
             r.run();
@@ -938,6 +947,7 @@ public void testSetJobState_withHandler_retriesUntilSuccess() throws IOException
         JobTask jobTask = mock(JobTask.class);
         AtomicInteger attempts = new AtomicInteger();
         doAnswer(invocation -> {
+            // Simulate transient failures for the first two attempts, then succeed on the third
             @SuppressWarnings("unchecked")
             ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>> listener = (ActionListener<
                 PersistentTasksCustomMetadata.PersistentTask<?>>) invocation.getArguments()[1];
@@ -954,18 +964,22 @@ public void testSetJobState_withHandler_retriesUntilSuccess() throws IOException
 
         manager.setJobState(jobTask, JobState.OPENED, "retry-test", handler);
 
+        // confirms that the method was called exactly three times (two failures then one success).
         verify(jobTask, times(3)).updatePersistentTaskState(any(JobTaskState.class), any());
         assertNull(holder.get());
     }
 
     public void testSetJobState_withHandler_noRetryOnResourceNotFound() throws IOException {
+        // Ensures that if the persistent‐state update fails with a ResourceNotFoundException, the retry loop does not retry
+        // again but immediately invokes the user’s handler with that exception.
         AutodetectProcessManager manager = createSpyManager();
         JobTask jobTask = mock(JobTask.class);
         ResourceNotFoundException rnfe = new ResourceNotFoundException("not found");
         doAnswer(invocation -> {
+            // Simulate a ResourceNotFoundException that should not be retried
             @SuppressWarnings("unchecked")
-            ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>> listener =
-                (ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>>) invocation.getArguments()[1];
+            ActionListener<PersistentTasksCustomMetadata.PersistentTask<?>> listener = (ActionListener<
+                PersistentTasksCustomMetadata.PersistentTask<?>>) invocation.getArguments()[1];
             listener.onFailure(rnfe);
             return null;
         }).when(jobTask).updatePersistentTaskState(any(), any());
@@ -975,7 +989,9 @@ public void testSetJobState_withHandler_noRetryOnResourceNotFound() throws IOExc
 
         manager.setJobState(jobTask, JobState.OPENED, "rnfe-test", handler);
 
+        // updatePersistentTaskState(...) was invoked exactly once (no retries).
         verify(jobTask, times(1)).updatePersistentTaskState(any(JobTaskState.class), any());
+        // The handler should have been invoked with the ResourceNotFoundException
         assertSame(rnfe, holder.get());
     }