test jvm update

Author: Andyz26

gradle.properties

@@ -1,4 +1,5 @@
 org.gradle.parallel=true
 org.gradle.caching=false
 org.gradle.configureondemand=true
-org.gradle.jvmargs=-Xmx1G "-XX:MaxMetaspaceSize=384m"
+org.gradle.jvmargs=-Xmx2G "-XX:MaxMetaspaceSize=512m"
+testJvmMaxHeap=2g

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

@@ -17,7 +17,9 @@
 package io.mantisrx.master.resourcecluster;
 
 import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertThrows;
 
 import akka.actor.ActorRef;
@@ -62,6 +64,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.stream.Collectors;
 import org.junit.AfterClass;
@@ -157,6 +160,7 @@ public static void teardown() {
     @Before
     public void setupRpcService() {
         rpcService.registerGateway(TASK_EXECUTOR_ADDRESS, gateway);
+        when(gateway.submitTask(any())).thenReturn(CompletableFuture.completedFuture(Ack.getInstance()));
         mantisJobStore = mock(MantisJobStore.class);
         jobMessageRouter = mock(JobMessageRouter.class);
     }

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

@@ -199,6 +199,7 @@ public void setupTest() throws Exception {
 
     private void setupRpcService() throws Exception {
         rpcService.registerGateway(TASK_EXECUTOR_ADDRESS, gateway);
+        when(gateway.submitTask(ArgumentMatchers.any())).thenReturn(CompletableFuture.completedFuture(Ack.getInstance()));
         mantisJobStore = mock(MantisJobStore.class);
         jobMessageRouter = mock(JobMessageRouter.class);
         doReturn(ImmutableList.of())
@@ -528,6 +529,11 @@ public void testGetTaskExecutorsUsageAndList() throws Exception {
 
     @Test
     public void testAssignmentTimeout() throws Exception {
+        // Override the default submitTask mock to return a hanging future that never completes
+        // This simulates a task executor that times out during assignment
+        CompletableFuture<Ack> hangingFuture = new CompletableFuture<>();
+        when(gateway.submitTask(ArgumentMatchers.any())).thenReturn(hangingFuture);
+
         assertEquals(Ack.getInstance(), resourceCluster.registerTaskExecutor(TASK_EXECUTOR_REGISTRATION).get());
         assertEquals(Ack.getInstance(),
             resourceCluster
@@ -546,6 +552,9 @@ public void testAssignmentTimeout() throws Exception {
         assertEquals(ImmutableList.of(), resourceCluster.getRegisteredTaskExecutors().get());
         assertEquals(ImmutableList.of(), resourceCluster.getAvailableTaskExecutors().get());
 
+        // Restore the submitTask mock to return a completed future for re-registration
+        when(gateway.submitTask(ArgumentMatchers.any())).thenReturn(CompletableFuture.completedFuture(Ack.getInstance()));
+
         // re-register TE
         when(mantisJobStore.getTaskExecutor(TASK_EXECUTOR_ID))
             .thenReturn(new TaskExecutorRegistration(TASK_EXECUTOR_ID, CLUSTER_ID, "", "", null, MACHINE_DEFINITION, ATTRIBUTES));