Fix a string operation for custom health check and update test (#1924)

diff --git a/helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/MaintenanceManagementService.java b/helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/MaintenanceManagementService.java
index 082c54b..c70b9c0 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/MaintenanceManagementService.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/MaintenanceManagementService.java
@@ -548,9 +548,9 @@
         // "CUSTOM_PARTITION_HEALTH_FAILURE:PARTITION_INITIAL_STATE_FAIL:partition_name"
         // we want to keep the first 2 parts as failed test name.
         String[] checks = failedCheck.split(":", 3);
-        failedCheck = checks[0] + checks[1];
+        failedCheck = checks[0] + ":" + checks[1];
       }
-      // Helix own health check name wil be in this pattern "HELIX:INSTANCE_NOT_ALIVE",
+      // Helix own health check name will be in this pattern "HELIX:INSTANCE_NOT_ALIVE",
       // no need to preprocess.
       if (!_nonBlockingHealthChecks.contains(failedCheck)) {
         return false;
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/clusterMaintenanceService/TestMaintenanceManagementService.java b/helix-rest/src/test/java/org/apache/helix/rest/clusterMaintenanceService/TestMaintenanceManagementService.java
index b24ffa1..a911286 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/clusterMaintenanceService/TestMaintenanceManagementService.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/clusterMaintenanceService/TestMaintenanceManagementService.java
@@ -29,6 +29,7 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import org.apache.helix.AccessOption;
 import org.apache.helix.BaseDataAccessor;
 import org.apache.helix.ConfigAccessor;
@@ -104,6 +105,7 @@
     }
   }
 
+
   @Test
   public void testGetInstanceStoppableCheckWhenHelixOwnCheckFail() throws IOException {
     Map<String, Boolean> failedCheck = ImmutableMap.of("FailCheck", false);
@@ -233,8 +235,28 @@
             false, HelixRestNamespace.DEFAULT_NAMESPACE_NAME);
     stoppableCheck = instanceServiceWithoutReadZK.getInstanceStoppableCheck(TEST_CLUSTER, TEST_INSTANCE, jsonContent);
     Assert.assertFalse(stoppableCheck.isStoppable());
+
+    // Test take instance with same setting.
+    MaintenanceManagementInstanceInfo instanceInfo =
+        instanceServiceWithoutReadZK.takeInstance(TEST_CLUSTER, TEST_INSTANCE, Collections.singletonList("CustomInstanceStoppableCheck"),
+            MaintenanceManagementService.getMapFromJsonPayload(jsonContent), Collections.singletonList("org.apache.helix.rest.server.TestOperationImpl"),
+            Collections.EMPTY_MAP, true);
+    Assert.assertFalse(instanceInfo.isSuccessful());
+    Assert.assertEquals(instanceInfo.getMessages().get(0), "CUSTOM_PARTITION_HEALTH_FAILURE:UNHEALTHY_PARTITION:PARTITION_0");
+
+    // Operation should finish even with check failed.
+    MockMaintenanceManagementService instanceServiceSkipFailure =
+        new MockMaintenanceManagementService(zkHelixDataAccessor, _configAccessor, _customRestClient, true,
+            ImmutableSet.of("CUSTOM_PARTITION_HEALTH_FAILURE:UNHEALTHY_PARTITION"), HelixRestNamespace.DEFAULT_NAMESPACE_NAME);
+    MaintenanceManagementInstanceInfo instanceInfo2 =
+        instanceServiceSkipFailure.takeInstance(TEST_CLUSTER, TEST_INSTANCE, Collections.singletonList("CustomInstanceStoppableCheck"),
+            MaintenanceManagementService.getMapFromJsonPayload(jsonContent), Collections.singletonList("org.apache.helix.rest.server.TestOperationImpl"),
+            Collections.EMPTY_MAP, true);
+    Assert.assertTrue(instanceInfo2.isSuccessful());
+    Assert.assertEquals(instanceInfo2.getOperationResult(), "DummyTakeOperationResult");
   }
-  /*
+
+    /*
    * Tests stoppable check api when all checks query is enabled. After helix own check fails,
    * the subsequent checks should be performed.
    */