TolerationConstraint implements ConstraintEvaluator { @Override public Result evaluate(TaskRequest taskRequest, VirtualMachineCurrentState targetVM, TaskTrackerState taskTrackerState) { Optional<AgentInstance> instanceOpt = SchedulerUtils.findInstance(agentManagementService, configuration.getInstanceAttributeName(), targetVM); if (!instanceOpt.isPresent()) { return MACHINE_DOES_NOT_EXIST; } AgentInstance agentInstance = instanceOpt.get(); Optional<AgentInstanceGroup> instanceGroupOpt = agentManagementService.findInstanceGroup(agentInstance.getInstanceGroupId()); if (!instanceGroupOpt.isPresent()) { return MACHINE_GROUP_DOES_NOT_EXIST; } AgentInstanceGroup instanceGroup = instanceGroupOpt.get(); Set<String> taints = getTaints(instanceGroup, agentInstance); return taints.contains(toleration) ? VALID : TOLERATION_DOES_NOT_MATCH; } TolerationConstraint(SchedulerConfiguration configuration, AgentManagementService agentManagementService, String machineId); @Override String getName(); @Override Result evaluate(TaskRequest taskRequest, VirtualMachineCurrentState targetVM, TaskTrackerState taskTrackerState); static final String NAME; }
@Test public void machineDoesNotExist() { when(agentManagementService.findAgentInstance(MACHINE_ID)).thenReturn(Optional.empty()); Result result = constraint.evaluate(createTaskRequest(TASK_ID), createVirtualMachineCurrentStateMock(MACHINE_ID), createTaskTrackerState()); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("The machine does not exist"); }
@Test public void machineGroupDoesExist() { AgentInstance instance = AgentGenerator.agentInstances().getValue().toBuilder().withId(MACHINE_ID).withInstanceGroupId("").build(); when(agentManagementService.findAgentInstance(MACHINE_ID)).thenReturn(Optional.of(instance)); Result result = constraint.evaluate(createTaskRequest(TASK_ID), createVirtualMachineCurrentStateMock(MACHINE_ID), createTaskTrackerState()); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("The machine group does not exist"); }
@Test public void machineTaintMatchesToleration() { AgentInstance instance = AgentGenerator.agentInstances().getValue().toBuilder().withId(MACHINE_ID) .withInstanceGroupId(MACHINE_GROUP) .withAttributes(Collections.singletonMap(TAINTS, TOLERATION)) .build(); when(agentManagementService.findAgentInstance(MACHINE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup instanceGroup = AgentGenerator.agentServerGroups().getValue().toBuilder() .withId(MACHINE_GROUP) .build(); when(agentManagementService.findInstanceGroup(MACHINE_GROUP)).thenReturn(Optional.of(instanceGroup)); Result result = constraint.evaluate(createTaskRequest(TASK_ID), createVirtualMachineCurrentStateMock(MACHINE_ID), createTaskTrackerState()); assertThat(result.isSuccessful()).isTrue(); }
@Test public void machineGroupTaintMatchesToleration() { AgentInstance instance = AgentGenerator.agentInstances().getValue().toBuilder().withId(MACHINE_ID) .withInstanceGroupId(MACHINE_GROUP) .build(); when(agentManagementService.findAgentInstance(MACHINE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup instanceGroup = AgentGenerator.agentServerGroups().getValue().toBuilder() .withId(MACHINE_GROUP) .withAttributes(Collections.singletonMap(TAINTS, TOLERATION)) .build(); when(agentManagementService.findInstanceGroup(MACHINE_GROUP)).thenReturn(Optional.of(instanceGroup)); Result result = constraint.evaluate(createTaskRequest(TASK_ID), createVirtualMachineCurrentStateMock(MACHINE_ID), createTaskTrackerState()); assertThat(result.isSuccessful()).isTrue(); }
@Test public void taintsDoNotMatchToleration() { AgentInstance instance = AgentGenerator.agentInstances().getValue().toBuilder().withId(MACHINE_ID) .withInstanceGroupId(MACHINE_GROUP) .build(); when(agentManagementService.findAgentInstance(MACHINE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup instanceGroup = AgentGenerator.agentServerGroups().getValue().toBuilder() .withId(MACHINE_GROUP) .build(); when(agentManagementService.findInstanceGroup(MACHINE_GROUP)).thenReturn(Optional.of(instanceGroup)); Result result = constraint.evaluate(createTaskRequest(TASK_ID), createVirtualMachineCurrentStateMock(MACHINE_ID), createTaskTrackerState()); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("The machine or machine group does not have a matching taint"); }