helix-core/src/test/java/org/apache/helix/task/assigner/TestAssignableInstance.java - helix - Git at Google

 package org.apache.helix.task.assigner;

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 import java.util.HashMap;
 import java.util.Map;

 import org.apache.helix.model.ClusterConfig;
 import org.apache.helix.model.InstanceConfig;
 import org.apache.helix.model.LiveInstance;
 import org.apache.helix.task.TaskConfig;
 import org.apache.helix.task.TaskConstants;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 import org.testng.collections.Maps;

 public class TestAssignableInstance extends AssignerTestBase {

   @Test
   public void testInvalidInitialization() {
     try {
       AssignableInstance ai = new AssignableInstance(null, null, null);
       Assert.fail("Expecting IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       Assert.assertTrue(e.getMessage().contains("cannot be null"));
     }

     try {
       ClusterConfig clusterConfig = new ClusterConfig("testCluster");
       InstanceConfig instanceConfig = new InstanceConfig("instance");
       LiveInstance liveInstance = new LiveInstance("another-instance");
       AssignableInstance ai = new AssignableInstance(clusterConfig, instanceConfig, liveInstance);
       Assert.fail("Expecting IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       Assert.assertTrue(e.getMessage().contains("don't match"));
     }
   }

   @Test
   public void testInitializationWithQuotaUnset() {
     int expectedCurrentTaskThreadPoolSize = 100;
     LiveInstance liveInstance = createLiveInstance(null, null);
     liveInstance.setCurrentTaskThreadPoolSize(expectedCurrentTaskThreadPoolSize);

     // Initialize AssignableInstance with neither resource capacity nor quota ratio provided
     AssignableInstance ai = new AssignableInstance(createClusterConfig(null, null, false),
         new InstanceConfig(testInstanceName), liveInstance);
     Assert.assertEquals(ai.getUsedCapacity().size(), 1);
     Assert.assertEquals(
         (int) ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
             .get(AssignableInstance.DEFAULT_QUOTA_TYPE), 0);
     Assert.assertEquals(
         (int) ai.getTotalCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
             .get(AssignableInstance.DEFAULT_QUOTA_TYPE), expectedCurrentTaskThreadPoolSize);
     Assert.assertEquals(ai.getCurrentAssignments().size(), 0);
   }

   @Test
   public void testInitializationWithOnlyCapacity() {
     // Initialize AssignableInstance with only resource capacity provided
     AssignableInstance ai = new AssignableInstance(createClusterConfig(null, null, false),
         new InstanceConfig(testInstanceName),
         createLiveInstance(testResourceTypes, testResourceCapacity));
     Assert.assertEquals(ai.getTotalCapacity().size(), testResourceTypes.length);
     Assert.assertEquals(ai.getUsedCapacity().size(), testResourceTypes.length);
     for (int i = 0; i < testResourceTypes.length; i++) {
       Assert.assertEquals(ai.getTotalCapacity().get(testResourceTypes[i]).size(), 1);
       Assert.assertEquals(ai.getUsedCapacity().get(testResourceTypes[i]).size(), 1);
       Assert.assertEquals(ai.getTotalCapacity().get(testResourceTypes[i])
           .get(AssignableInstance.DEFAULT_QUOTA_TYPE), Integer.valueOf(testResourceCapacity[i]));
       Assert.assertEquals(
           ai.getUsedCapacity().get(testResourceTypes[i]).get(AssignableInstance.DEFAULT_QUOTA_TYPE),
           Integer.valueOf(0));
     }
   }

   @Test
   public void testInitializationWithOnlyQuotaType() {
     int expectedCurrentTaskThreadPoolSize = 100;
     LiveInstance liveInstance = createLiveInstance(null, null);
     liveInstance.setCurrentTaskThreadPoolSize(expectedCurrentTaskThreadPoolSize);

     // Initialize AssignableInstance with only quota type provided
     AssignableInstance ai =
         new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
             new InstanceConfig(testInstanceName), liveInstance);

     Assert.assertEquals(ai.getTotalCapacity().size(), 1);
     Assert.assertEquals(ai.getUsedCapacity().size(), 1);
     Assert.assertEquals(
         ai.getTotalCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()).size(),
         testQuotaTypes.length);
     Assert.assertEquals(
         ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()).size(),
         testQuotaTypes.length);
     Assert.assertEquals(
         ai.getTotalCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()),
         calculateExpectedQuotaPerType(expectedCurrentTaskThreadPoolSize, testQuotaTypes,
             testQuotaRatio));
     Assert.assertEquals(ai.getCurrentAssignments().size(), 0);
   }

   @Test
   public void testInitializationWithQuotaAndCapacity() {
     // Initialize AssignableInstance with both capacity and quota type provided
     AssignableInstance ai =
         new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
             new InstanceConfig(testInstanceName),
             createLiveInstance(testResourceTypes, testResourceCapacity));

     Map<String, Integer> usedResourcePerType =
         createResourceQuotaPerTypeMap(testQuotaTypes, new int[] {
             0, 0, 0
         });
     for (int i = 0; i < testResourceTypes.length; i++) {
       Assert.assertEquals(ai.getTotalCapacity().get(testResourceTypes[i]),
           calculateExpectedQuotaPerType(Integer.valueOf(testResourceCapacity[i]), testQuotaTypes,
               testQuotaRatio));
       Assert.assertEquals(ai.getUsedCapacity().get(testResourceTypes[i]), usedResourcePerType);
     }
   }

   @Test
   public void testAssignableInstanceUpdateConfigs() {
     AssignableInstance ai =
         new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
             new InstanceConfig(testInstanceName),
             createLiveInstance(testResourceTypes, testResourceCapacity));

     String[] newResources = new String[] {
         "Resource2", "Resource3", "Resource4"
     };
     String[] newResourceCapacities = new String[] {
         "100", "150", "50"
     };

     String[] newTypes = new String[] {
         "Type3", "Type4", "Type5", "Type6"
     };
     String[] newTypeRatio = new String[] {
         "20", "40", "25", "25"
     };

     LiveInstance newLiveInstance = createLiveInstance(newResources, newResourceCapacities);
     ClusterConfig newClusterConfig = createClusterConfig(newTypes, newTypeRatio, false);
     ai.updateConfigs(newClusterConfig, null, newLiveInstance);

     Assert.assertEquals(ai.getUsedCapacity().size(), newResourceCapacities.length);
     Assert.assertEquals(ai.getTotalCapacity().size(), newResourceCapacities.length);

     for (int i = 0; i < newResources.length; i++) {
       Assert.assertEquals(ai.getTotalCapacity().get(newResources[i]), calculateExpectedQuotaPerType(
           Integer.valueOf(newResourceCapacities[i]), newTypes, newTypeRatio));
       Assert.assertEquals(ai.getUsedCapacity().get(newResources[i]),
           createResourceQuotaPerTypeMap(newTypes, new int[] {
               0, 0, 0, 0
           }));
     }
   }

   @Test
   public void testNormalTryAssign() {
     int testCurrentTaskThreadPoolSize = 100;
     LiveInstance liveInstance = createLiveInstance(null, null);
     liveInstance.setCurrentTaskThreadPoolSize(testCurrentTaskThreadPoolSize);

     AssignableInstance ai = new AssignableInstance(createClusterConfig(null, null, true),
         new InstanceConfig(testInstanceName), liveInstance);

     // When nothing is configured, we should use default quota type to assign
     Map<String, TaskAssignResult> results = new HashMap<>();
     for (int i = 0; i < testCurrentTaskThreadPoolSize; i++) {
       String taskId = Integer.toString(i);
       TaskConfig task = new TaskConfig("", null, taskId, null);
       TaskAssignResult result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
       Assert.assertTrue(result.isSuccessful());
       ai.assign(result);
       results.put(taskId, result);
     }

     // We are out of quota now and we should not be able to assign
     String taskId = "TaskCannotAssign";
     TaskConfig task = new TaskConfig("", null, taskId, null);
     TaskAssignResult result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
     Assert.assertFalse(result.isSuccessful());
     Assert.assertEquals(result.getFailureReason(),
         TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
     try {
       ai.assign(result);
       Assert.fail("Expecting IllegalStateException");
     } catch (IllegalStateException e) {
       // OK
     }

     // After releasing 1 task, we should be able to schedule
     ai.release(results.get("1").getTaskConfig(), AssignableInstance.DEFAULT_QUOTA_TYPE);
     result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
     Assert.assertTrue(result.isSuccessful());

     // release all tasks, check remaining resources
     for (TaskAssignResult rst : results.values()) {
       ai.release(rst.getTaskConfig(), AssignableInstance.DEFAULT_QUOTA_TYPE);
     }

     Assert.assertEquals(
         (int) ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
             .get(AssignableInstance.DEFAULT_QUOTA_TYPE),
         0);
   }

   @Test
   public void testTryAssignFailure() {
     AssignableInstance ai =
         new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
             new InstanceConfig(testInstanceName),
             createLiveInstance(testResourceTypes, testResourceCapacity));

     // No such resource type
     String taskId = "testTask";
     TaskConfig task = new TaskConfig("", null, taskId, "");
     TaskAssignResult result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
     Assert.assertFalse(result.isSuccessful());
     Assert.assertEquals(result.getFailureReason(),
         TaskAssignResult.FailureReason.NO_SUCH_RESOURCE_TYPE);

     ai.updateConfigs(null, null, createLiveInstance(new String[] {
         LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()
     }, new String[] {
         "1"
     }));

     ai.updateConfigs(createClusterConfig(testQuotaTypes, testQuotaRatio, true), null, null);

     result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);

     Assert.assertTrue(result.isSuccessful());
     ai.assign(result);
     try {
       ai.assign(result);
       Assert.fail("Expecting IllegalArgumentException");
     } catch (IllegalStateException e) {
       // OK
     }

     // Duplicate assignment
     result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
     Assert.assertFalse(result.isSuccessful());
     Assert.assertEquals(result.getFailureReason(),
         TaskAssignResult.FailureReason.TASK_ALREADY_ASSIGNED);

     // Insufficient quota
     ai.release(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
     ai.updateConfigs(null, null, createLiveInstance(new String[] {
         LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()
     }, new String[] {
         "0"
     }));

     result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);

     Assert.assertFalse(result.isSuccessful());
     Assert.assertEquals(result.getFailureReason(),
         TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
   }

   @Test
   public void testRestoreTaskAssignResult() {
     AssignableInstance ai =
         new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, true),
             new InstanceConfig(testInstanceName), createLiveInstance(new String[] {
                 LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()
             }, new String[] {
                 "40"
             }));

     Map<String, TaskConfig> currentAssignments = new HashMap<>();
     TaskConfig supportedTask = new TaskConfig("", null, "supportedTask", "");
     currentAssignments.put("supportedTask", supportedTask);

     TaskConfig unsupportedTask = new TaskConfig("", null, "unsupportedTask", "");
     currentAssignments.put("unsupportedTask", unsupportedTask);

     Map<String, TaskAssignResult> results = Maps.newHashMap();
     for (Map.Entry<String, TaskConfig> entry : currentAssignments.entrySet()) {
       String taskID = entry.getKey();
       TaskConfig taskConfig = entry.getValue();
       String quotaType = (taskID.equals("supportedTask")) ? AssignableInstance.DEFAULT_QUOTA_TYPE
           : "UnsupportedQuotaType";
       // Restore TaskAssignResult
       TaskAssignResult taskAssignResult = ai.restoreTaskAssignResult(taskID, taskConfig, quotaType);
       if (taskAssignResult.isSuccessful()) {
         results.put(taskID, taskAssignResult);
       }
     }

     for (TaskAssignResult rst : results.values()) {
       Assert.assertTrue(rst.isSuccessful());
       Assert.assertEquals(rst.getAssignableInstance(), ai);
     }
     Assert.assertEquals(ai.getCurrentAssignments().size(), 2);
     // The expected value for the following should be 2, not 1 because the unsupported task should
     // also have been assigned as a DEFAULT task
     Assert.assertEquals(
         (int) ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
             .get(AssignableInstance.DEFAULT_QUOTA_TYPE),
         2);
   }

   private Map<String, Integer> createResourceQuotaPerTypeMap(String[] types, int[] quotas) {
     Map<String, Integer> ret = new HashMap<>();
     for (int i = 0; i < types.length; i++) {
       ret.put(types[i], quotas[i]);
     }
     return ret;
   }

   private Map<String, Integer> calculateExpectedQuotaPerType(int capacity, String[] quotaTypes,
       String[] quotaRatios) {
     Integer totalQuota = 0;
     Map<String, Integer> expectedQuotaPerType = new HashMap<>();

     for (String ratio : quotaRatios) {
       totalQuota += Integer.valueOf(ratio);
     }

     for (int i = 0; i < quotaRatios.length; i++) {
       expectedQuotaPerType.put(quotaTypes[i],
           Math.round((float) capacity * Integer.valueOf(quotaRatios[i]) / totalQuota));
     }
     return expectedQuotaPerType;
   }
 }
	package org.apache.helix.task.assigner;

	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	import java.util.HashMap;
	import java.util.Map;

	import org.apache.helix.model.ClusterConfig;
	import org.apache.helix.model.InstanceConfig;
	import org.apache.helix.model.LiveInstance;
	import org.apache.helix.task.TaskConfig;
	import org.apache.helix.task.TaskConstants;
	import org.testng.Assert;
	import org.testng.annotations.Test;
	import org.testng.collections.Maps;

	public class TestAssignableInstance extends AssignerTestBase {

	@Test
	public void testInvalidInitialization() {
	try {
	AssignableInstance ai = new AssignableInstance(null, null, null);
	Assert.fail("Expecting IllegalArgumentException");
	} catch (IllegalArgumentException e) {
	Assert.assertTrue(e.getMessage().contains("cannot be null"));
	}

	try {
	ClusterConfig clusterConfig = new ClusterConfig("testCluster");
	InstanceConfig instanceConfig = new InstanceConfig("instance");
	LiveInstance liveInstance = new LiveInstance("another-instance");
	AssignableInstance ai = new AssignableInstance(clusterConfig, instanceConfig, liveInstance);
	Assert.fail("Expecting IllegalArgumentException");
	} catch (IllegalArgumentException e) {
	Assert.assertTrue(e.getMessage().contains("don't match"));
	}
	}

	@Test
	public void testInitializationWithQuotaUnset() {
	int expectedCurrentTaskThreadPoolSize = 100;
	LiveInstance liveInstance = createLiveInstance(null, null);
	liveInstance.setCurrentTaskThreadPoolSize(expectedCurrentTaskThreadPoolSize);

	// Initialize AssignableInstance with neither resource capacity nor quota ratio provided
	AssignableInstance ai = new AssignableInstance(createClusterConfig(null, null, false),
	new InstanceConfig(testInstanceName), liveInstance);
	Assert.assertEquals(ai.getUsedCapacity().size(), 1);
	Assert.assertEquals(
	(int) ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
	.get(AssignableInstance.DEFAULT_QUOTA_TYPE), 0);
	Assert.assertEquals(
	(int) ai.getTotalCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
	.get(AssignableInstance.DEFAULT_QUOTA_TYPE), expectedCurrentTaskThreadPoolSize);
	Assert.assertEquals(ai.getCurrentAssignments().size(), 0);
	}

	@Test
	public void testInitializationWithOnlyCapacity() {
	// Initialize AssignableInstance with only resource capacity provided
	AssignableInstance ai = new AssignableInstance(createClusterConfig(null, null, false),
	new InstanceConfig(testInstanceName),
	createLiveInstance(testResourceTypes, testResourceCapacity));
	Assert.assertEquals(ai.getTotalCapacity().size(), testResourceTypes.length);
	Assert.assertEquals(ai.getUsedCapacity().size(), testResourceTypes.length);
	for (int i = 0; i < testResourceTypes.length; i++) {
	Assert.assertEquals(ai.getTotalCapacity().get(testResourceTypes[i]).size(), 1);
	Assert.assertEquals(ai.getUsedCapacity().get(testResourceTypes[i]).size(), 1);
	Assert.assertEquals(ai.getTotalCapacity().get(testResourceTypes[i])
	.get(AssignableInstance.DEFAULT_QUOTA_TYPE), Integer.valueOf(testResourceCapacity[i]));
	Assert.assertEquals(
	ai.getUsedCapacity().get(testResourceTypes[i]).get(AssignableInstance.DEFAULT_QUOTA_TYPE),
	Integer.valueOf(0));
	}
	}

	@Test
	public void testInitializationWithOnlyQuotaType() {
	int expectedCurrentTaskThreadPoolSize = 100;
	LiveInstance liveInstance = createLiveInstance(null, null);
	liveInstance.setCurrentTaskThreadPoolSize(expectedCurrentTaskThreadPoolSize);

	// Initialize AssignableInstance with only quota type provided
	AssignableInstance ai =
	new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
	new InstanceConfig(testInstanceName), liveInstance);

	Assert.assertEquals(ai.getTotalCapacity().size(), 1);
	Assert.assertEquals(ai.getUsedCapacity().size(), 1);
	Assert.assertEquals(
	ai.getTotalCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()).size(),
	testQuotaTypes.length);
	Assert.assertEquals(
	ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()).size(),
	testQuotaTypes.length);
	Assert.assertEquals(
	ai.getTotalCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()),
	calculateExpectedQuotaPerType(expectedCurrentTaskThreadPoolSize, testQuotaTypes,
	testQuotaRatio));
	Assert.assertEquals(ai.getCurrentAssignments().size(), 0);
	}

	@Test
	public void testInitializationWithQuotaAndCapacity() {
	// Initialize AssignableInstance with both capacity and quota type provided
	AssignableInstance ai =
	new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
	new InstanceConfig(testInstanceName),
	createLiveInstance(testResourceTypes, testResourceCapacity));

	Map<String, Integer> usedResourcePerType =
	createResourceQuotaPerTypeMap(testQuotaTypes, new int[] {
	0, 0, 0
	});
	for (int i = 0; i < testResourceTypes.length; i++) {
	Assert.assertEquals(ai.getTotalCapacity().get(testResourceTypes[i]),
	calculateExpectedQuotaPerType(Integer.valueOf(testResourceCapacity[i]), testQuotaTypes,
	testQuotaRatio));
	Assert.assertEquals(ai.getUsedCapacity().get(testResourceTypes[i]), usedResourcePerType);
	}
	}

	@Test
	public void testAssignableInstanceUpdateConfigs() {
	AssignableInstance ai =
	new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
	new InstanceConfig(testInstanceName),
	createLiveInstance(testResourceTypes, testResourceCapacity));

	String[] newResources = new String[] {
	"Resource2", "Resource3", "Resource4"
	};
	String[] newResourceCapacities = new String[] {
	"100", "150", "50"
	};

	String[] newTypes = new String[] {
	"Type3", "Type4", "Type5", "Type6"
	};
	String[] newTypeRatio = new String[] {
	"20", "40", "25", "25"
	};

	LiveInstance newLiveInstance = createLiveInstance(newResources, newResourceCapacities);
	ClusterConfig newClusterConfig = createClusterConfig(newTypes, newTypeRatio, false);
	ai.updateConfigs(newClusterConfig, null, newLiveInstance);

	Assert.assertEquals(ai.getUsedCapacity().size(), newResourceCapacities.length);
	Assert.assertEquals(ai.getTotalCapacity().size(), newResourceCapacities.length);

	for (int i = 0; i < newResources.length; i++) {
	Assert.assertEquals(ai.getTotalCapacity().get(newResources[i]), calculateExpectedQuotaPerType(
	Integer.valueOf(newResourceCapacities[i]), newTypes, newTypeRatio));
	Assert.assertEquals(ai.getUsedCapacity().get(newResources[i]),
	createResourceQuotaPerTypeMap(newTypes, new int[] {
	0, 0, 0, 0
	}));
	}
	}

	@Test
	public void testNormalTryAssign() {
	int testCurrentTaskThreadPoolSize = 100;
	LiveInstance liveInstance = createLiveInstance(null, null);
	liveInstance.setCurrentTaskThreadPoolSize(testCurrentTaskThreadPoolSize);

	AssignableInstance ai = new AssignableInstance(createClusterConfig(null, null, true),
	new InstanceConfig(testInstanceName), liveInstance);

	// When nothing is configured, we should use default quota type to assign
	Map<String, TaskAssignResult> results = new HashMap<>();
	for (int i = 0; i < testCurrentTaskThreadPoolSize; i++) {
	String taskId = Integer.toString(i);
	TaskConfig task = new TaskConfig("", null, taskId, null);
	TaskAssignResult result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
	Assert.assertTrue(result.isSuccessful());
	ai.assign(result);
	results.put(taskId, result);
	}

	// We are out of quota now and we should not be able to assign
	String taskId = "TaskCannotAssign";
	TaskConfig task = new TaskConfig("", null, taskId, null);
	TaskAssignResult result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
	Assert.assertFalse(result.isSuccessful());
	Assert.assertEquals(result.getFailureReason(),
	TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
	try {
	ai.assign(result);
	Assert.fail("Expecting IllegalStateException");
	} catch (IllegalStateException e) {
	// OK
	}

	// After releasing 1 task, we should be able to schedule
	ai.release(results.get("1").getTaskConfig(), AssignableInstance.DEFAULT_QUOTA_TYPE);
	result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
	Assert.assertTrue(result.isSuccessful());

	// release all tasks, check remaining resources
	for (TaskAssignResult rst : results.values()) {
	ai.release(rst.getTaskConfig(), AssignableInstance.DEFAULT_QUOTA_TYPE);
	}

	Assert.assertEquals(
	(int) ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
	.get(AssignableInstance.DEFAULT_QUOTA_TYPE),
	0);
	}

	@Test
	public void testTryAssignFailure() {
	AssignableInstance ai =
	new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, false),
	new InstanceConfig(testInstanceName),
	createLiveInstance(testResourceTypes, testResourceCapacity));

	// No such resource type
	String taskId = "testTask";
	TaskConfig task = new TaskConfig("", null, taskId, "");
	TaskAssignResult result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
	Assert.assertFalse(result.isSuccessful());
	Assert.assertEquals(result.getFailureReason(),
	TaskAssignResult.FailureReason.NO_SUCH_RESOURCE_TYPE);

	ai.updateConfigs(null, null, createLiveInstance(new String[] {
	LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()
	}, new String[] {
	"1"
	}));

	ai.updateConfigs(createClusterConfig(testQuotaTypes, testQuotaRatio, true), null, null);

	result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);

	Assert.assertTrue(result.isSuccessful());
	ai.assign(result);
	try {
	ai.assign(result);
	Assert.fail("Expecting IllegalArgumentException");
	} catch (IllegalStateException e) {
	// OK
	}

	// Duplicate assignment
	result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
	Assert.assertFalse(result.isSuccessful());
	Assert.assertEquals(result.getFailureReason(),
	TaskAssignResult.FailureReason.TASK_ALREADY_ASSIGNED);

	// Insufficient quota
	ai.release(task, AssignableInstance.DEFAULT_QUOTA_TYPE);
	ai.updateConfigs(null, null, createLiveInstance(new String[] {
	LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()
	}, new String[] {
	"0"
	}));

	result = ai.tryAssign(task, AssignableInstance.DEFAULT_QUOTA_TYPE);

	Assert.assertFalse(result.isSuccessful());
	Assert.assertEquals(result.getFailureReason(),
	TaskAssignResult.FailureReason.INSUFFICIENT_QUOTA);
	}

	@Test
	public void testRestoreTaskAssignResult() {
	AssignableInstance ai =
	new AssignableInstance(createClusterConfig(testQuotaTypes, testQuotaRatio, true),
	new InstanceConfig(testInstanceName), createLiveInstance(new String[] {
	LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name()
	}, new String[] {
	"40"
	}));

	Map<String, TaskConfig> currentAssignments = new HashMap<>();
	TaskConfig supportedTask = new TaskConfig("", null, "supportedTask", "");
	currentAssignments.put("supportedTask", supportedTask);

	TaskConfig unsupportedTask = new TaskConfig("", null, "unsupportedTask", "");
	currentAssignments.put("unsupportedTask", unsupportedTask);

	Map<String, TaskAssignResult> results = Maps.newHashMap();
	for (Map.Entry<String, TaskConfig> entry : currentAssignments.entrySet()) {
	String taskID = entry.getKey();
	TaskConfig taskConfig = entry.getValue();
	String quotaType = (taskID.equals("supportedTask")) ? AssignableInstance.DEFAULT_QUOTA_TYPE
	: "UnsupportedQuotaType";
	// Restore TaskAssignResult
	TaskAssignResult taskAssignResult = ai.restoreTaskAssignResult(taskID, taskConfig, quotaType);
	if (taskAssignResult.isSuccessful()) {
	results.put(taskID, taskAssignResult);
	}
	}

	for (TaskAssignResult rst : results.values()) {
	Assert.assertTrue(rst.isSuccessful());
	Assert.assertEquals(rst.getAssignableInstance(), ai);
	}
	Assert.assertEquals(ai.getCurrentAssignments().size(), 2);
	// The expected value for the following should be 2, not 1 because the unsupported task should
	// also have been assigned as a DEFAULT task
	Assert.assertEquals(
	(int) ai.getUsedCapacity().get(LiveInstance.InstanceResourceType.TASK_EXEC_THREAD.name())
	.get(AssignableInstance.DEFAULT_QUOTA_TYPE),
	2);
	}

	private Map<String, Integer> createResourceQuotaPerTypeMap(String[] types, int[] quotas) {
	Map<String, Integer> ret = new HashMap<>();
	for (int i = 0; i < types.length; i++) {
	ret.put(types[i], quotas[i]);
	}
	return ret;
	}

	private Map<String, Integer> calculateExpectedQuotaPerType(int capacity, String[] quotaTypes,
	String[] quotaRatios) {
	Integer totalQuota = 0;
	Map<String, Integer> expectedQuotaPerType = new HashMap<>();

	for (String ratio : quotaRatios) {
	totalQuota += Integer.valueOf(ratio);
	}

	for (int i = 0; i < quotaRatios.length; i++) {
	expectedQuotaPerType.put(quotaTypes[i],
	Math.round((float) capacity * Integer.valueOf(quotaRatios[i]) / totalQuota));
	}
	return expectedQuotaPerType;
	}
	}