diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestPartitionAssignmentAPI.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestPartitionAssignmentAPI.java
index f3a4bbdd9b..cab348c21d 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestPartitionAssignmentAPI.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestPartitionAssignmentAPI.java
@@ -103,8 +103,6 @@ public void beforeTest() {
         .setResources(new HashSet<>(_resources))
         .setWaitTillVerify(TestHelper.DEFAULT_REBALANCE_PROCESSING_WAIT_TIME).build();
 
-    Assert.assertTrue(_clusterVerifier.verifyByPolling());
-
     // Add and start instances to cluster
     for (int i = 0; i < DEFAULT_INSTANCE_COUNT; i++) {
       String instanceName = INSTANCE_NAME_PREFIX + (INSTANCE_START_PORT + i);
@@ -119,6 +117,7 @@ public void beforeTest() {
       _participants.add(participant);
     }
 
+    Assert.assertTrue(_clusterVerifier.verifyByPolling());
     System.out.println("End setup:" + TestHelper.getTestMethodName());
   }
 
@@ -126,11 +125,8 @@ public void beforeTest() {
   public void afterTest() throws Exception {
     System.out.println("Start teardown:" + TestHelper.getTestMethodName());
 
-    // Drop all resources
-    for (String resource : _resources) {
-      _gSetupTool.dropResourceFromCluster(CLUSTER_NAME, resource);
-    }
-    _resources.clear();
+    // Stop controller
+    _controller.syncStop();
 
     // Stop and remove all instances
     for (MockParticipantManager participant : _participants) {
@@ -143,8 +139,11 @@ public void afterTest() throws Exception {
     }
     _participants.clear();
 
-    // Stop controller
-    _controller.syncStop();
+    // Drop all resources
+    for (String resource : _resources) {
+      _gSetupTool.dropResourceFromCluster(CLUSTER_NAME, resource);
+    }
+    _resources.clear();
 
     // Drop cluster
     _gSetupTool.deleteCluster(CLUSTER_NAME);
@@ -442,6 +441,7 @@ private void testComputePartitionAssignmentMaintenanceMode() throws Exception {
     MockParticipantManager toAddParticipant = createParticipant(toAddInstanceName);
     toAddParticipant.syncStart();
 
+    Assert.assertTrue(_clusterVerifier.verifyByPolling());
     // Choose participant to simulate killing in API call
     MockParticipantManager participantToKill = _participants.get(0);
 
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestPerInstanceAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestPerInstanceAccessor.java
index 58dbb78ea7..8ba32e0ccd 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/server/TestPerInstanceAccessor.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/server/TestPerInstanceAccessor.java
@@ -23,6 +23,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -41,24 +42,43 @@
 import org.apache.helix.HelixException;
 import org.apache.helix.TestHelper;
 import org.apache.helix.constants.InstanceConstants;
+import org.apache.helix.integration.manager.MockParticipantManager;
+import org.apache.helix.integration.task.MockTask;
 import org.apache.helix.manager.zk.ZKHelixDataAccessor;
 import org.apache.helix.model.ClusterConfig;
 import org.apache.helix.model.ExternalView;
-import org.apache.helix.model.IdealState;
 import org.apache.helix.model.InstanceConfig;
 import org.apache.helix.model.Message;
+import org.apache.helix.participant.StateMachineEngine;
 import org.apache.helix.rest.server.resources.AbstractResource;
 import org.apache.helix.rest.server.resources.helix.InstancesAccessor;
 import org.apache.helix.rest.server.resources.helix.PerInstanceAccessor;
 import org.apache.helix.rest.server.util.JerseyUriRequestBuilder;
+import org.apache.helix.task.TaskFactory;
+import org.apache.helix.task.TaskStateModelFactory;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 public class TestPerInstanceAccessor extends AbstractTestClass {
   private final static String CLUSTER_NAME = "TestCluster_4";
   private final static String INSTANCE_NAME = CLUSTER_NAME + "localhost_12918";
 
+  private MockParticipantManager _instanceToDisable;
+
+  @BeforeClass
+  public void beforeClass() {
+    int indexToDisable = -1;
+    for (int i = 0; i < _mockParticipantManagers.size(); i++) {
+      if (_mockParticipantManagers.get(i).getInstanceName().equals(INSTANCE_NAME)) {
+        indexToDisable = i;
+        break;
+      }
+    }
+    _instanceToDisable = _mockParticipantManagers.remove(indexToDisable);
+  }
+
   @Test
   public void testIsInstanceStoppable() throws IOException {
     System.out.println("Start test :" + TestHelper.getTestMethodName());
@@ -254,10 +274,11 @@ public void testTakeInstanceCheckOnly() throws IOException {
   }
 
   @Test(dependsOnMethods = "testTakeInstanceCheckOnly")
-  public void testGetAllMessages() throws IOException {
+  public void testGetAllMessages() throws Exception {
     System.out.println("Start test :" + TestHelper.getTestMethodName());
-    _mockParticipantManagers.get(0).disconnect();
-    String testInstance = CLUSTER_NAME + "localhost_12918"; //Non-live instance
+    _instanceToDisable.disconnect();
+
+    String testInstance = INSTANCE_NAME; //Non-live instance
 
     String messageId = "msg1";
     Message message = new Message(Message.MessageType.STATE_TRANSITION, messageId);
@@ -277,14 +298,14 @@ public void testGetAllMessages() throws IOException {
         node.get(PerInstanceAccessor.PerInstanceProperties.total_message_count.name()).intValue();
 
     Assert.assertEquals(newMessageCount, 1);
+    helixDataAccessor.removeProperty(helixDataAccessor.keyBuilder().message(testInstance, messageId));
     System.out.println("End test :" + TestHelper.getTestMethodName());
   }
 
   @Test(dependsOnMethods = "testGetAllMessages")
   public void testGetMessagesByStateModelDef() throws Exception {
     System.out.println("Start test :" + TestHelper.getTestMethodName());
-
-    String testInstance = CLUSTER_NAME + "localhost_12918"; //Non-live instance
+    String testInstance = INSTANCE_NAME; //Non-live instance
     String messageId = "msg1";
     Message message = new Message(Message.MessageType.STATE_TRANSITION, messageId);
     message.setStateModelDef("MasterSlave");
@@ -315,7 +336,15 @@ public void testGetMessagesByStateModelDef() throws Exception {
         node.get(PerInstanceAccessor.PerInstanceProperties.total_message_count.name()).intValue();
 
     Assert.assertEquals(newMessageCount, 0);
-    _mockParticipantManagers.get(0).connect();
+    MockParticipantManager participant =
+        new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, INSTANCE_NAME);
+    Map<String, TaskFactory> taskFactoryReg = new HashMap<>();
+    taskFactoryReg.put(MockTask.TASK_COMMAND, MockTask::new);
+    StateMachineEngine stateMachineEngine = participant.getStateMachineEngine();
+    stateMachineEngine.registerStateModelFactory("Task",
+        new TaskStateModelFactory(participant, taskFactoryReg));
+    participant.syncStart();
+    _mockParticipantManagers.add(participant);
     System.out.println("End test :" + TestHelper.getTestMethodName());
   }