test fix

Author: Andyz26

mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/ResourceClusterScalerActor.java

@@ -309,7 +309,7 @@ private void onGetClusterIdleInstancesResponse(GetClusterIdleInstancesResponse r
     private void onTriggerClusterUsageRequest(TriggerClusterUsageRequest req) {
         log.trace("Requesting cluster usage: {}", this.clusterId);
         if (this.skuToRuleMap.isEmpty()) {
-            log.info("{} scaler is disabled due to no rules", this.clusterId);
+            log.debug("{} scaler is disabled due to no rules", this.clusterId);
             return;
         }
 

mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/TaskExecutorState.java

@@ -20,7 +20,10 @@
 import io.mantisrx.master.resourcecluster.ResourceClusterActor.AvailabilityState;
 import io.mantisrx.master.resourcecluster.ResourceClusterActor.Pending;
 import io.mantisrx.master.resourcecluster.ResourceClusterActor.Running;
+import io.mantisrx.master.jobcluster.job.worker.WorkerState;
+import io.mantisrx.master.jobcluster.job.worker.WorkerTerminate;
 import io.mantisrx.server.core.domain.WorkerId;
+import io.mantisrx.server.core.JobCompletedReason;
 import io.mantisrx.server.master.resourcecluster.TaskExecutorHeartbeat;
 import io.mantisrx.server.master.resourcecluster.TaskExecutorRegistration;
 import io.mantisrx.server.master.resourcecluster.TaskExecutorReport;
@@ -199,6 +202,14 @@ boolean onHeartbeat(TaskExecutorHeartbeat heartbeat)
         }
 
         TaskExecutorReport report = heartbeat.getTaskExecutorReport();
+        if (this.availabilityState instanceof Running && report instanceof Available) {
+            WorkerId runningWorkerId = this.availabilityState.getWorkerId();
+            if (runningWorkerId != null) {
+                log.warn("Heartbeat indicates available while running {}. Marking worker as lost.", runningWorkerId);
+                jobMessageRouter.routeWorkerEvent(
+                    new WorkerTerminate(runningWorkerId, WorkerState.Failed, JobCompletedReason.Lost));
+            }
+        }
         if (this.cancelledWorkerOnTask != null) {
             if (report instanceof Occupied && ((Occupied) report).getWorkerId().equals(this.cancelledWorkerOnTask)) {
                 log.warn("{} cancelled, request cancel on heartbeat.", this.cancelledWorkerOnTask);

mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ResourceClusterActorTest.java

@@ -269,9 +269,8 @@ public void testMarkTaskCancelled() throws Exception {
         doReturn(ImmutableList.of())
             .when(mantisJobStore)
             .getJobArtifactsToCache(CLUSTER_ID);
-        doReturn(TASK_EXECUTOR_REGISTRATION)
-            .when(mantisJobStore)
-            .getTaskExecutor(TASK_EXECUTOR_ID);
+        when(mantisJobStore.getTaskExecutor(TASK_EXECUTOR_ID))
+            .thenReturn(TASK_EXECUTOR_REGISTRATION);
 
         assertEquals(
             Ack.getInstance(),
@@ -316,9 +315,8 @@ public void testInitializationAfterRestart() throws Exception {
         doReturn(ImmutableList.of())
             .when(mantisJobStore)
             .getJobArtifactsToCache(CLUSTER_ID);
-        doReturn(TASK_EXECUTOR_REGISTRATION)
-            .when(mantisJobStore)
-            .getTaskExecutor(TASK_EXECUTOR_ID);
+        when(mantisJobStore.getTaskExecutor(ArgumentMatchers.any(TaskExecutorID.class)))
+            .thenReturn(TASK_EXECUTOR_REGISTRATION);
         assertEquals(
             Ack.getInstance(),
             resourceCluster.initializeTaskExecutor(TASK_EXECUTOR_ID, WORKER_ID).get());
@@ -549,9 +547,8 @@ public void testAssignmentTimeout() throws Exception {
         assertEquals(ImmutableList.of(), resourceCluster.getAvailableTaskExecutors().get());
 
         // re-register TE
-        doReturn(new TaskExecutorRegistration(TASK_EXECUTOR_ID, CLUSTER_ID, "", "", null, MACHINE_DEFINITION, ATTRIBUTES))
-            .when(mantisJobStore)
-            .getTaskExecutor(TASK_EXECUTOR_ID);
+        when(mantisJobStore.getTaskExecutor(TASK_EXECUTOR_ID))
+            .thenReturn(new TaskExecutorRegistration(TASK_EXECUTOR_ID, CLUSTER_ID, "", "", null, MACHINE_DEFINITION, ATTRIBUTES));
         assertEquals(Ack.getInstance(),
             resourceCluster
                 .heartBeatFromTaskExecutor(
@@ -802,7 +799,7 @@ public void testGetAssignedTaskExecutorAfterTaskCompletes() throws Throwable {
 
     @Test
     public void testTaskExecutorIsDisabledEvenAfterRestart() throws Exception {
-        doReturn(TASK_EXECUTOR_REGISTRATION).when(mantisJobStore).getTaskExecutor(TASK_EXECUTOR_ID);
+        when(mantisJobStore.getTaskExecutor(TASK_EXECUTOR_ID)).thenReturn(TASK_EXECUTOR_REGISTRATION);
         doReturn(ImmutableList.of()).when(mantisJobStore).getJobArtifactsToCache(CLUSTER_ID);
 
         resourceCluster.registerTaskExecutor(TASK_EXECUTOR_REGISTRATION);