Repository: aurora Updated Branches: refs/heads/master 5b34231ba -> 4e6242fed
http://git-wip-us.apache.org/repos/asf/aurora/blob/4e6242fe/src/test/java/org/apache/aurora/scheduler/preemptor/PreemptionVictimFilterTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/preemptor/PreemptionVictimFilterTest.java b/src/test/java/org/apache/aurora/scheduler/preemptor/PreemptionVictimFilterTest.java index 348386e..b3ffb0d 100644 --- a/src/test/java/org/apache/aurora/scheduler/preemptor/PreemptionVictimFilterTest.java +++ b/src/test/java/org/apache/aurora/scheduler/preemptor/PreemptionVictimFilterTest.java @@ -39,13 +39,13 @@ import org.apache.aurora.gen.ScheduledTask; import org.apache.aurora.gen.TaskConfig; import org.apache.aurora.gen.TaskEvent; import org.apache.aurora.gen.apiConstants; -import org.apache.aurora.scheduler.TierInfo; -import org.apache.aurora.scheduler.TierManager; +import org.apache.aurora.scheduler.base.TaskTestUtil; import org.apache.aurora.scheduler.filter.SchedulingFilter; import org.apache.aurora.scheduler.filter.SchedulingFilter.Veto; import org.apache.aurora.scheduler.filter.SchedulingFilterImpl; import org.apache.aurora.scheduler.mesos.TaskExecutors; import org.apache.aurora.scheduler.offers.HostOffer; +import org.apache.aurora.scheduler.preemptor.PreemptionVictimFilter.PreemptionVictimFilterImpl; import org.apache.aurora.scheduler.resources.ResourceBag; import org.apache.aurora.scheduler.resources.ResourceTestUtil; import org.apache.aurora.scheduler.resources.ResourceType; @@ -67,9 +67,6 @@ import static org.apache.aurora.gen.MaintenanceMode.NONE; import static org.apache.aurora.gen.Resource.numCpus; import static org.apache.aurora.gen.Resource.ramMb; import static org.apache.aurora.gen.ScheduleStatus.RUNNING; -import static org.apache.aurora.scheduler.base.TaskTestUtil.DEV_TIER; -import static org.apache.aurora.scheduler.base.TaskTestUtil.PREFERRED_TIER; -import static org.apache.aurora.scheduler.base.TaskTestUtil.REVOCABLE_TIER; import static org.apache.aurora.scheduler.filter.AttributeAggregate.empty; import static org.apache.aurora.scheduler.preemptor.PreemptionVictimFilter.PreemptionVictimFilterImpl.ORDER; import static org.apache.aurora.scheduler.preemptor.PreemptorMetrics.MISSING_ATTRIBUTES_NAME; @@ -110,7 +107,6 @@ public class PreemptionVictimFilterTest extends EasyMockTest { private SchedulingFilter schedulingFilter; private FakeStatsProvider statsProvider; private PreemptorMetrics preemptorMetrics; - private TierManager tierManager; private FakeClock clock; @Before @@ -119,7 +115,6 @@ public class PreemptionVictimFilterTest extends EasyMockTest { storageUtil.expectOperations(); statsProvider = new FakeStatsProvider(); preemptorMetrics = new PreemptorMetrics(new CachedCounters(statsProvider)); - tierManager = createMock(TierManager.class); clock = new FakeClock(); ResourceType.initializeEmptyCliArgsForTest(); } @@ -129,12 +124,11 @@ public class PreemptionVictimFilterTest extends EasyMockTest { Optional<HostOffer> offer, ScheduledTask... victims) { - PreemptionVictimFilter.PreemptionVictimFilterImpl filter = - new PreemptionVictimFilter.PreemptionVictimFilterImpl( - schedulingFilter, - TaskExecutors.NO_OVERHEAD_EXECUTOR, - preemptorMetrics, - tierManager); + PreemptionVictimFilterImpl filter = new PreemptionVictimFilterImpl( + schedulingFilter, + TaskExecutors.NO_OVERHEAD_EXECUTOR, + preemptorMetrics, + TaskTestUtil.TIER_MANAGER); return filter.filterPreemptionVictims( ITaskConfig.build(pendingTask.getAssignedTask().getTask()), @@ -145,16 +139,14 @@ public class PreemptionVictimFilterTest extends EasyMockTest { } @Test - public void testPreempted() throws Exception { + public void testPreempted() { setUpHost(); schedulingFilter = createMock(SchedulingFilter.class); ScheduledTask lowPriority = makeTask(USER_A, JOB_A, TASK_ID_A); assignToHost(lowPriority); - expectGetTier(lowPriority, DEV_TIER).times(2); ScheduledTask highPriority = makeTask(USER_A, JOB_A, TASK_ID_B, 100); - expectGetTier(highPriority, DEV_TIER); expectFiltering(); @@ -163,7 +155,7 @@ public class PreemptionVictimFilterTest extends EasyMockTest { } @Test - public void testLowestPriorityPreempted() throws Exception { + public void testLowestPriorityPreempted() { setUpHost(); schedulingFilter = createMock(SchedulingFilter.class); @@ -172,10 +164,8 @@ public class PreemptionVictimFilterTest extends EasyMockTest { ScheduledTask lowerPriority = makeTask(USER_A, JOB_A, TASK_ID_B, 1); assignToHost(lowerPriority); - expectGetTier(lowerPriority, DEV_TIER).atLeastOnce(); ScheduledTask highPriority = makeTask(USER_A, JOB_A, TASK_ID_C, 100); - expectGetTier(highPriority, DEV_TIER); expectFiltering(); @@ -184,24 +174,20 @@ public class PreemptionVictimFilterTest extends EasyMockTest { } @Test - public void testOnePreemptableTask() throws Exception { + public void testOnePreemptableTask() { setUpHost(); schedulingFilter = createMock(SchedulingFilter.class); ScheduledTask highPriority = makeTask(USER_A, JOB_A, TASK_ID_A, 100); assignToHost(highPriority); - expectGetTier(highPriority, DEV_TIER); ScheduledTask lowerPriority = makeTask(USER_A, JOB_A, TASK_ID_B, 99); assignToHost(lowerPriority); - expectGetTier(lowerPriority, DEV_TIER); ScheduledTask lowestPriority = makeTask(USER_A, JOB_A, TASK_ID_C, 1); assignToHost(lowestPriority); - expectGetTier(lowestPriority, DEV_TIER).times(2); ScheduledTask pendingPriority = makeTask(USER_A, JOB_A, TASK_ID_D, 98); - expectGetTier(pendingPriority, DEV_TIER).times(3); expectFiltering(); @@ -212,29 +198,25 @@ public class PreemptionVictimFilterTest extends EasyMockTest { } @Test - public void testHigherPriorityRunning() throws Exception { + public void testHigherPriorityRunning() { schedulingFilter = createMock(SchedulingFilter.class); ScheduledTask highPriority = makeTask(USER_A, JOB_A, TASK_ID_B, 100); assignToHost(highPriority); - expectGetTier(highPriority, DEV_TIER); ScheduledTask task = makeTask(USER_A, JOB_A, TASK_ID_A); - expectGetTier(task, DEV_TIER); control.replay(); assertNoVictims(runFilter(task, NO_OFFER, highPriority)); } @Test - public void testProductionPreemptingNonproduction() throws Exception { + public void testProductionPreemptingNonproduction() { setUpHost(); schedulingFilter = createMock(SchedulingFilter.class); // Use a very low priority for the production task to show that priority is irrelevant. ScheduledTask p1 = makeProductionTask(USER_A, JOB_A, TASK_ID_A + "_p1", -1000); - expectGetTier(p1, PREFERRED_TIER); ScheduledTask a1 = makeTask(USER_A, JOB_A, TASK_ID_B + "_a1", 100); - expectGetTier(a1, DEV_TIER).times(2); assignToHost(a1); expectFiltering(); @@ -244,16 +226,14 @@ public class PreemptionVictimFilterTest extends EasyMockTest { } @Test - public void testProductionPreemptingNonproductionAcrossUsers() throws Exception { + public void testProductionPreemptingNonproductionAcrossUsers() { setUpHost(); schedulingFilter = createMock(SchedulingFilter.class); // Use a very low priority for the production task to show that priority is irrelevant. ScheduledTask p1 = makeProductionTask(USER_A, JOB_A, TASK_ID_A + "_p1", -1000); - expectGetTier(p1, PREFERRED_TIER); ScheduledTask a1 = makeTask(USER_B, JOB_A, TASK_ID_B + "_a1", 100); assignToHost(a1); - expectGetTier(a1, DEV_TIER).times(2); expectFiltering(); @@ -262,12 +242,10 @@ public class PreemptionVictimFilterTest extends EasyMockTest { } @Test - public void testProductionUsersDoNotPreemptEachOther() throws Exception { + public void testProductionUsersDoNotPreemptEachOther() { schedulingFilter = createMock(SchedulingFilter.class); ScheduledTask p1 = makeProductionTask(USER_A, JOB_A, TASK_ID_A + "_p1", 1000); - expectGetTier(p1, PREFERRED_TIER); ScheduledTask a1 = makeProductionTask(USER_B, JOB_A, TASK_ID_B + "_a1", 0); - expectGetTier(a1, PREFERRED_TIER); assignToHost(a1); control.replay(); @@ -276,17 +254,15 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures a production task can preempt 2 tasks on the same host. @Test - public void testProductionPreemptingManyNonProduction() throws Exception { + public void testProductionPreemptingManyNonProduction() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); ScheduledTask a1 = makeTask(USER_A, JOB_A, TASK_ID_A + "_a1"); setResource(a1, CPUS, 1.0); setResource(a1, RAM_MB, 512.0); - expectGetTier(a1, DEV_TIER).atLeastOnce(); ScheduledTask b1 = makeTask(USER_B, JOB_B, TASK_ID_B + "_b1"); setResource(b1, CPUS, 1.0); setResource(b1, RAM_MB, 512.0); - expectGetTier(b1, DEV_TIER).atLeastOnce(); setUpHost(); @@ -296,7 +272,6 @@ public class PreemptionVictimFilterTest extends EasyMockTest { ScheduledTask p1 = makeProductionTask(USER_B, JOB_B, TASK_ID_B + "_p1"); setResource(p1, CPUS, 2.0); setResource(p1, RAM_MB, 1024.0); - expectGetTier(p1, PREFERRED_TIER).times(2); control.replay(); assertVictims(runFilter(p1, NO_OFFER, a1, b1), a1, b1); @@ -304,12 +279,11 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures we select the minimal number of tasks to preempt @Test - public void testMinimalSetPreempted() throws Exception { + public void testMinimalSetPreempted() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); ScheduledTask a1 = makeTask(USER_A, JOB_A, TASK_ID_A + "_a1"); setResource(a1, CPUS, 4.0); setResource(a1, RAM_MB, 4096.0); - expectGetTier(a1, DEV_TIER).atLeastOnce(); ScheduledTask b1 = makeTask(USER_B, JOB_B, TASK_ID_B + "_b1"); b1.getAssignedTask().getTask() @@ -318,12 +292,10 @@ public class PreemptionVictimFilterTest extends EasyMockTest { ramMb(512))); setResource(b1, CPUS, 1.0); setResource(b1, RAM_MB, 512.0); - expectGetTier(b1, DEV_TIER).anyTimes(); ScheduledTask b2 = makeTask(USER_B, JOB_B, TASK_ID_B + "_b2"); setResource(b2, CPUS, 1.0); setResource(b2, RAM_MB, 512.0); - expectGetTier(b2, DEV_TIER).anyTimes(); setUpHost(); @@ -334,7 +306,6 @@ public class PreemptionVictimFilterTest extends EasyMockTest { ScheduledTask p1 = makeProductionTask(USER_C, JOB_C, TASK_ID_C + "_p1"); setResource(p1, CPUS, 2.0); setResource(p1, RAM_MB, 1024.0); - expectGetTier(p1, PREFERRED_TIER).times(3); control.replay(); assertVictims(runFilter(p1, NO_OFFER, b1, b2, a1), a1); @@ -342,14 +313,13 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures a production task *never* preempts a production task from another job. @Test - public void testProductionJobNeverPreemptsProductionJob() throws Exception { + public void testProductionJobNeverPreemptsProductionJob() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); ScheduledTask p1 = makeProductionTask(USER_A, JOB_A, TASK_ID_A + "_p1"); p1.getAssignedTask().getTask() .setResources(ImmutableSet.of( numCpus(2), ramMb(1024))); - expectGetTier(p1, PREFERRED_TIER); setUpHost(); @@ -360,7 +330,6 @@ public class PreemptionVictimFilterTest extends EasyMockTest { .setResources(ImmutableSet.of( numCpus(1), ramMb(512))); - expectGetTier(p2, PREFERRED_TIER); control.replay(); assertNoVictims(runFilter(p2, NO_OFFER, p1)); @@ -368,7 +337,7 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures that we can preempt if a task + offer can satisfy a pending task. @Test - public void testPreemptWithOfferAndTask() throws Exception { + public void testPreemptWithOfferAndTask() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); setUpHost(); @@ -379,14 +348,12 @@ public class PreemptionVictimFilterTest extends EasyMockTest { numCpus(1), ramMb(512))); assignToHost(a1); - expectGetTier(a1, DEV_TIER).times(2); ScheduledTask p1 = makeProductionTask(USER_B, JOB_B, TASK_ID_B + "_p1"); p1.getAssignedTask().getTask() .setResources(ImmutableSet.of( numCpus(2), ramMb(1024))); - expectGetTier(p1, PREFERRED_TIER); control.replay(); assertVictims( @@ -399,7 +366,7 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures revocable offer resources are filtered out. @Test - public void testRevocableOfferFiltered() throws Exception { + public void testRevocableOfferFiltered() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); setUpHost(); @@ -408,12 +375,10 @@ public class PreemptionVictimFilterTest extends EasyMockTest { setResource(a1, CPUS, 1.0); setResource(a1, RAM_MB, 512.0); assignToHost(a1); - expectGetTier(a1, DEV_TIER).times(2); ScheduledTask p1 = makeProductionTask(USER_B, JOB_B, TASK_ID_B + "_p1"); setResource(p1, CPUS, 2.0); setResource(p1, RAM_MB, 1024.0); - expectGetTier(p1, PREFERRED_TIER); control.replay(); assertNoVictims(runFilter( @@ -424,21 +389,20 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures revocable task CPU is not considered for preemption. @Test - public void testRevocableVictimsFiltered() throws Exception { + public void testRevocableVictimsFiltered() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); setUpHost(); ScheduledTask a1 = makeTask(USER_A, JOB_A, TASK_ID_A + "_a1"); + a1.getAssignedTask().getTask().setTier(TaskTestUtil.REVOCABLE_TIER_NAME); setResource(a1, CPUS, 1.0); setResource(a1, RAM_MB, 512.0); assignToHost(a1); - expectGetTier(a1, REVOCABLE_TIER).times(2); ScheduledTask p1 = makeProductionTask(USER_B, JOB_B, TASK_ID_B + "_p1"); setResource(p1, CPUS, 2.0); setResource(p1, RAM_MB, 1024.0); - expectGetTier(p1, PREFERRED_TIER); control.replay(); assertNoVictims(runFilter( @@ -449,7 +413,7 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures revocable victim non-compressible resources are still considered. @Test - public void testRevocableVictimRamUsed() throws Exception { + public void testRevocableVictimRamUsed() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); setUpHost(); @@ -458,12 +422,10 @@ public class PreemptionVictimFilterTest extends EasyMockTest { setResource(a1, CPUS, 1.0); setResource(a1, RAM_MB, 512.0); assignToHost(a1); - expectGetTier(a1, REVOCABLE_TIER).times(2); ScheduledTask p1 = makeProductionTask(USER_B, JOB_B, TASK_ID_B + "_p1"); setResource(p1, CPUS, 2.0); setResource(p1, RAM_MB, 1024.0); - expectGetTier(p1, PREFERRED_TIER); control.replay(); assertVictims( @@ -476,7 +438,7 @@ public class PreemptionVictimFilterTest extends EasyMockTest { // Ensures we can preempt if two tasks and an offer can satisfy a pending task. @Test - public void testPreemptWithOfferAndMultipleTasks() throws Exception { + public void testPreemptWithOfferAndMultipleTasks() { schedulingFilter = new SchedulingFilterImpl(UNAVAILABLITY_THRESHOLD, clock); setUpHost(); @@ -485,18 +447,15 @@ public class PreemptionVictimFilterTest extends EasyMockTest { setResource(a1, CPUS, 1.0); setResource(a1, RAM_MB, 512.0); assignToHost(a1); - expectGetTier(a1, DEV_TIER).atLeastOnce(); ScheduledTask a2 = makeTask(USER_A, JOB_B, TASK_ID_A + "_a2"); setResource(a2, CPUS, 1.0); setResource(a2, RAM_MB, 512.0); assignToHost(a2); - expectGetTier(a2, DEV_TIER).atLeastOnce(); ScheduledTask p1 = makeProductionTask(USER_B, JOB_B, TASK_ID_B + "_p1"); setResource(p1, CPUS, 4.0); setResource(p1, RAM_MB, 2048.0); - expectGetTier(p1, PREFERRED_TIER).times(2); control.replay(); Optional<HostOffer> offer = @@ -519,13 +478,11 @@ public class PreemptionVictimFilterTest extends EasyMockTest { schedulingFilter = createMock(SchedulingFilter.class); ScheduledTask task = makeProductionTask(USER_A, JOB_A, TASK_ID_A); assignToHost(task); - expectGetTier(task, PREFERRED_TIER); ScheduledTask a1 = makeTask(USER_A, JOB_A, TASK_ID_A + "_a1"); setResource(a1, CPUS, 1.0); setResource(a1, RAM_MB, 512.0); assignToHost(a1); - expectGetTier(a1, DEV_TIER); expect(storageUtil.attributeStore.getHostAttributes(HOST_A)).andReturn(Optional.empty()); @@ -540,13 +497,11 @@ public class PreemptionVictimFilterTest extends EasyMockTest { schedulingFilter = createMock(SchedulingFilter.class); ScheduledTask task = makeProductionTask(USER_A, JOB_A, TASK_ID_A); assignToHost(task); - expectGetTier(task, PREFERRED_TIER); ScheduledTask a1 = makeTask(USER_A, JOB_A, TASK_ID_A + "_a1"); setResource(a1, CPUS, 1.0); setResource(a1, RAM_MB, 512.0); assignToHost(a1); - expectGetTier(a1, DEV_TIER).times(2); setUpHost(); expectFiltering(Optional.of(Veto.constraintMismatch("ban"))); @@ -661,11 +616,6 @@ public class PreemptionVictimFilterTest extends EasyMockTest { .andAnswer(() -> veto.map(ImmutableSet::of).orElse(ImmutableSet.of())); } - private IExpectationSetters<TierInfo> expectGetTier(ScheduledTask task, TierInfo tier) { - return expect(tierManager.getTier(ITaskConfig.build(task.getAssignedTask().getTask()))) - .andReturn(tier); - } - private static void setResource(ScheduledTask task, ResourceType type, Double value) { task.getAssignedTask().setTask(ResourceTestUtil.resetResource( ITaskConfig.build(task.getAssignedTask().getTask()), @@ -686,17 +636,18 @@ public class PreemptionVictimFilterTest extends EasyMockTest { .setTask(new TaskConfig() .setJob(new JobKey(role, env, job)) .setPriority(priority) + .setTier(production ? TaskTestUtil.PROD_TIER_NAME : null) .setProduction(production) .setConstraints(Sets.newHashSet()) .setExecutorConfig(new ExecutorConfig(apiConstants.AURORA_EXECUTOR_NAME, "config"))); return new ScheduledTask().setAssignedTask(assignedTask); } - static ScheduledTask makeTask(String role, String job, String taskId) { + private static ScheduledTask makeTask(String role, String job, String taskId) { return makeTask(role, job, taskId, 0, "dev", false); } - static void addEvent(ScheduledTask task, ScheduleStatus status) { + private static void addEvent(ScheduledTask task, ScheduleStatus status) { task.addToTaskEvents(new TaskEvent(0, status)); } http://git-wip-us.apache.org/repos/asf/aurora/blob/4e6242fe/src/test/java/org/apache/aurora/scheduler/scheduling/FirstFitOfferSelectorTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/scheduling/FirstFitOfferSelectorTest.java b/src/test/java/org/apache/aurora/scheduler/scheduling/FirstFitOfferSelectorTest.java index 5e2fdcb..e6b2b74 100644 --- a/src/test/java/org/apache/aurora/scheduler/scheduling/FirstFitOfferSelectorTest.java +++ b/src/test/java/org/apache/aurora/scheduler/scheduling/FirstFitOfferSelectorTest.java @@ -16,15 +16,14 @@ package org.apache.aurora.scheduler.scheduling; import com.google.common.collect.ImmutableList; import org.apache.aurora.common.testing.easymock.EasyMockTest; +import org.apache.aurora.scheduler.base.TaskTestUtil; import org.apache.aurora.scheduler.offers.HostOffer; -import org.apache.aurora.scheduler.resources.ResourceBag; import org.apache.aurora.scheduler.storage.entities.IAssignedTask; import org.junit.Before; import org.junit.Test; import static org.apache.aurora.scheduler.base.TaskTestUtil.JOB; import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask; -import static org.apache.aurora.scheduler.filter.AttributeAggregate.empty; import static org.apache.aurora.scheduler.filter.SchedulingFilter.ResourceRequest; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -32,10 +31,8 @@ import static org.junit.Assert.assertFalse; public class FirstFitOfferSelectorTest extends EasyMockTest { private static final IAssignedTask TASK = makeTask("id", JOB).getAssignedTask(); - private static final ResourceRequest EMPTY_REQUEST = new ResourceRequest( - TASK.getTask(), - ResourceBag.EMPTY, - empty()); + private static final ResourceRequest EMPTY_REQUEST = + TaskTestUtil.toResourceRequest(TASK.getTask()); private OfferSelector firstFitOfferSelector; http://git-wip-us.apache.org/repos/asf/aurora/blob/4e6242fe/src/test/java/org/apache/aurora/scheduler/scheduling/TaskAssignerImplTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/scheduling/TaskAssignerImplTest.java b/src/test/java/org/apache/aurora/scheduler/scheduling/TaskAssignerImplTest.java index 78e1269..f84941e 100644 --- a/src/test/java/org/apache/aurora/scheduler/scheduling/TaskAssignerImplTest.java +++ b/src/test/java/org/apache/aurora/scheduler/scheduling/TaskAssignerImplTest.java @@ -26,28 +26,22 @@ import org.apache.aurora.gen.Attribute; import org.apache.aurora.gen.HostAttributes; import org.apache.aurora.gen.JobKey; import org.apache.aurora.gen.TaskConfig; -import org.apache.aurora.scheduler.TierManager; import org.apache.aurora.scheduler.base.InstanceKeys; import org.apache.aurora.scheduler.base.TaskGroupKey; -import org.apache.aurora.scheduler.base.Tasks; +import org.apache.aurora.scheduler.base.TaskTestUtil; import org.apache.aurora.scheduler.filter.AttributeAggregate; import org.apache.aurora.scheduler.filter.SchedulingFilter.ResourceRequest; import org.apache.aurora.scheduler.mesos.MesosTaskFactory; import org.apache.aurora.scheduler.offers.HostOffer; import org.apache.aurora.scheduler.offers.OfferManager; -import org.apache.aurora.scheduler.resources.ResourceBag; import org.apache.aurora.scheduler.state.StateChangeResult; import org.apache.aurora.scheduler.state.StateManager; import org.apache.aurora.scheduler.storage.entities.IAssignedTask; import org.apache.aurora.scheduler.storage.entities.IHostAttributes; import org.apache.aurora.scheduler.storage.entities.IInstanceKey; -import org.apache.aurora.scheduler.storage.entities.IScheduledTask; import org.apache.aurora.scheduler.storage.entities.ITaskConfig; import org.apache.aurora.scheduler.testing.FakeStatsProvider; import org.apache.aurora.scheduler.updater.UpdateAgentReserver; -import org.apache.mesos.v1.Protos.AgentID; -import org.apache.mesos.v1.Protos.FrameworkID; -import org.apache.mesos.v1.Protos.OfferID; import org.apache.mesos.v1.Protos.TaskID; import org.apache.mesos.v1.Protos.TaskInfo; import org.junit.Before; @@ -55,10 +49,10 @@ import org.junit.Test; import static org.apache.aurora.gen.ScheduleStatus.LOST; import static org.apache.aurora.gen.ScheduleStatus.PENDING; -import static org.apache.aurora.scheduler.base.TaskTestUtil.DEV_TIER; import static org.apache.aurora.scheduler.base.TaskTestUtil.JOB; import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask; import static org.apache.aurora.scheduler.filter.AttributeAggregate.empty; +import static org.apache.aurora.scheduler.mesos.TaskExecutors.NO_OVERHEAD_EXECUTOR; import static org.apache.aurora.scheduler.resources.ResourceTestUtil.mesosRange; import static org.apache.aurora.scheduler.resources.ResourceTestUtil.mesosScalar; import static org.apache.aurora.scheduler.resources.ResourceTestUtil.offer; @@ -88,27 +82,22 @@ public class TaskAssignerImplTest extends EasyMockTest { .setHost(MESOS_OFFER.getHostname()) .setAttributes(ImmutableSet.of( new Attribute("host", ImmutableSet.of(MESOS_OFFER.getHostname())))))); - private static final IScheduledTask TASK = makeTask("id", JOB); - private static final TaskGroupKey GROUP_KEY = TaskGroupKey.from(TASK.getAssignedTask().getTask()); + private static final IAssignedTask TASK = makeTask("id", JOB).getAssignedTask(); + private static final TaskGroupKey GROUP_KEY = TaskGroupKey.from(TASK.getTask()); private static final TaskInfo TASK_INFO = TaskInfo.newBuilder() .setName("taskName") - .setTaskId(TaskID.newBuilder().setValue(Tasks.id(TASK))) + .setTaskId(TaskID.newBuilder().setValue(TASK.getTaskId())) .setAgentId(MESOS_OFFER.getAgentId()) .build(); - private static final IInstanceKey INSTANCE_KEY = - InstanceKeys.from(JOB, TASK.getAssignedTask().getInstanceId()); + private static final IInstanceKey INSTANCE_KEY = InstanceKeys.from(JOB, TASK.getInstanceId()); private static final Map<String, TaskGroupKey> NO_RESERVATION = ImmutableMap.of(); - private static final HostOffer OFFER_2 = new HostOffer( - Offer.newBuilder() - .setId(OfferID.newBuilder().setValue("offerId0")) - .setFrameworkId(FrameworkID.newBuilder().setValue("frameworkId")) - .setAgentId(AgentID.newBuilder().setValue("slaveId0")) - .setHostname("hostName0") - .addResources(mesosRange(PORTS, PORT)) - .addResources(mesosScalar(CPUS, 1)) - .addResources(mesosScalar(RAM_MB, 1024)) - .build(), - IHostAttributes.build(new HostAttributes())); + private static final Offer MESOS_OFFER_2 = + offer("offer-2", mesosScalar(CPUS, 1), mesosScalar(RAM_MB, 1024), mesosRange(PORTS, PORT)); + private static final HostOffer OFFER_2 = + new HostOffer(MESOS_OFFER_2, IHostAttributes.build(new HostAttributes() + .setHost(MESOS_OFFER_2.getHostname()) + .setAttributes(ImmutableSet.of( + new Attribute("host", ImmutableSet.of(MESOS_OFFER_2.getHostname())))))); private static final Set<String> NO_ASSIGNMENT = ImmutableSet.of(); @@ -120,17 +109,15 @@ public class TaskAssignerImplTest extends EasyMockTest { private MesosTaskFactory taskFactory; private OfferManager offerManager; private TaskAssignerImpl assigner; - private TierManager tierManager; private FakeStatsProvider statsProvider; private UpdateAgentReserver updateAgentReserver; @Before - public void setUp() throws Exception { + public void setUp() { storeProvider = createMock(MutableStoreProvider.class); taskFactory = createMock(MesosTaskFactory.class); stateManager = createMock(StateManager.class); offerManager = createMock(OfferManager.class); - tierManager = createMock(TierManager.class); updateAgentReserver = createMock(UpdateAgentReserver.class); statsProvider = new FakeStatsProvider(); // TODO(jly): FirstFitOfferSelector returns the first offer which is what we want for testing, @@ -140,43 +127,49 @@ public class TaskAssignerImplTest extends EasyMockTest { stateManager, taskFactory, offerManager, - tierManager, updateAgentReserver, statsProvider, offerSelector); aggregate = empty(); - resourceRequest = new ResourceRequest( - TASK.getAssignedTask().getTask(), - ResourceBag.EMPTY, - aggregate); + resourceRequest = ResourceRequest.fromTask( + TASK.getTask(), + NO_OVERHEAD_EXECUTOR, + aggregate, + TaskTestUtil.TIER_MANAGER); } @Test - public void testAssignNoTasks() throws Exception { + public void testAssignNoTasks() { control.replay(); assertEquals( NO_ASSIGNMENT, - assigner.maybeAssign(storeProvider, null, null, ImmutableSet.of(), null)); + assigner.maybeAssign( + storeProvider, + resourceRequest, + GROUP_KEY, + ImmutableSet.of(), + NO_RESERVATION)); } @Test public void testAssignmentClearedOnError() throws Exception { - expectNoUpdateReservations(1); - expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest, false)) - .andReturn(ImmutableSet.of(OFFER, OFFER_2)); + expect(updateAgentReserver.isReserved(anyString())).andReturn(false).atLeastOnce(); + expect(updateAgentReserver.getAgent(anyObject())).andReturn(Optional.empty()).atLeastOnce(); + + expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest)) + .andReturn(ImmutableSet.of(OFFER, OFFER_2)).atLeastOnce(); offerManager.launchTask(MESOS_OFFER.getId(), TASK_INFO); expectLastCall().andThrow(new OfferManager.LaunchException("expected")); - expect(tierManager.getTier(TASK.getAssignedTask().getTask())).andReturn(DEV_TIER); expectAssignTask(MESOS_OFFER); expect(stateManager.changeState( storeProvider, - Tasks.id(TASK), + TASK.getTaskId(), Optional.of(PENDING), LOST, LAUNCH_FAILED_MSG)) .andReturn(StateChangeResult.SUCCESS); - expect(taskFactory.createFrom(TASK.getAssignedTask(), MESOS_OFFER, false)) + expect(taskFactory.createFrom(TASK, MESOS_OFFER, false)) .andReturn(TASK_INFO); control.replay(); @@ -188,9 +181,9 @@ public class TaskAssignerImplTest extends EasyMockTest { assigner.maybeAssign( storeProvider, resourceRequest, - TaskGroupKey.from(TASK.getAssignedTask().getTask()), + TaskGroupKey.from(TASK.getTask()), ImmutableSet.of( - TASK.getAssignedTask(), + TASK, makeTask("id2", JOB).getAssignedTask(), makeTask("id3", JOB).getAssignedTask()), NO_RESERVATION)); @@ -198,10 +191,9 @@ public class TaskAssignerImplTest extends EasyMockTest { } @Test - public void testAssignmentSkippedForReservedSlave() throws Exception { + public void testAssignmentSkippedForReservedSlave() { expectNoUpdateReservations(0); - expect(tierManager.getTier(TASK.getAssignedTask().getTask())).andReturn(DEV_TIER); - expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest, false)) + expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest)) .andReturn(ImmutableSet.of(OFFER)); control.replay(); @@ -211,8 +203,8 @@ public class TaskAssignerImplTest extends EasyMockTest { assigner.maybeAssign( storeProvider, resourceRequest, - TaskGroupKey.from(TASK.getAssignedTask().getTask()), - ImmutableSet.of(TASK.getAssignedTask()), + TaskGroupKey.from(TASK.getTask()), + ImmutableSet.of(TASK), ImmutableMap.of(SLAVE_ID, TaskGroupKey.from( ITaskConfig.build(new TaskConfig().setJob(new JobKey("other", "e", "n"))))))); } @@ -223,73 +215,68 @@ public class TaskAssignerImplTest extends EasyMockTest { // and permissive in task->slave direction. In other words, a task with a slave reservation // should still be tried against other unreserved slaves. expectNoUpdateReservations(1); - expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest, false)) + expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest)) .andReturn(ImmutableSet.of(OFFER_2, OFFER)); - expect(tierManager.getTier(TASK.getAssignedTask().getTask())).andReturn(DEV_TIER); expectAssignTask(OFFER_2.getOffer()); - expect(taskFactory.createFrom(TASK.getAssignedTask(), OFFER_2.getOffer(), false)) + expect(taskFactory.createFrom(TASK, OFFER_2.getOffer(), false)) .andReturn(TASK_INFO); offerManager.launchTask(OFFER_2.getOffer().getId(), TASK_INFO); control.replay(); assertEquals( - ImmutableSet.of(Tasks.id(TASK)), + ImmutableSet.of(TASK.getTaskId()), assigner.maybeAssign( storeProvider, resourceRequest, - TaskGroupKey.from(TASK.getAssignedTask().getTask()), - ImmutableSet.of(TASK.getAssignedTask()), + TaskGroupKey.from(TASK.getTask()), + ImmutableSet.of(TASK), ImmutableMap.of(SLAVE_ID, GROUP_KEY))); } @Test public void testResourceMapperCallback() { - AssignedTask builder = TASK.newBuilder().getAssignedTask(); + AssignedTask builder = TASK.newBuilder(); builder.unsetAssignedPorts(); control.replay(); assertEquals( - TASK.getAssignedTask(), + TASK, assigner.mapAndAssignResources(MESOS_OFFER, IAssignedTask.build(builder))); } @Test public void testAssignToReservedAgent() throws Exception { - expect(updateAgentReserver.hasReservations(GROUP_KEY)).andReturn(true); expect(updateAgentReserver.getAgent(INSTANCE_KEY)).andReturn(Optional.of(SLAVE_ID)); updateAgentReserver.release(SLAVE_ID, INSTANCE_KEY); - expect(offerManager.getMatching(MESOS_OFFER.getAgentId(), resourceRequest, false)) + expect(offerManager.getMatching(MESOS_OFFER.getAgentId(), resourceRequest)) .andReturn(Optional.of(OFFER)); expectAssignTask(MESOS_OFFER); offerManager.launchTask(MESOS_OFFER.getId(), TASK_INFO); - expect(tierManager.getTier(TASK.getAssignedTask().getTask())).andReturn(DEV_TIER); - expect(taskFactory.createFrom(TASK.getAssignedTask(), MESOS_OFFER, false)) + expect(taskFactory.createFrom(TASK, MESOS_OFFER, false)) .andReturn(TASK_INFO); control.replay(); assertEquals( - ImmutableSet.of(Tasks.id(TASK)), + ImmutableSet.of(TASK.getTaskId()), assigner.maybeAssign( storeProvider, resourceRequest, - TaskGroupKey.from(TASK.getAssignedTask().getTask()), + TaskGroupKey.from(TASK.getTask()), ImmutableSet.of( - TASK.getAssignedTask()), + TASK), ImmutableMap.of(SLAVE_ID, GROUP_KEY))); assertNotEquals(empty(), aggregate); } @Test - public void testAssignReservedAgentWhenOfferNotReady() throws Exception { - expect(updateAgentReserver.hasReservations(GROUP_KEY)).andReturn(true); + public void testAssignReservedAgentWhenOfferNotReady() { expect(updateAgentReserver.getAgent(INSTANCE_KEY)).andReturn(Optional.of(SLAVE_ID)); - expect(offerManager.getMatching(MESOS_OFFER.getAgentId(), resourceRequest, false)) + expect(offerManager.getMatching(MESOS_OFFER.getAgentId(), resourceRequest)) .andReturn(Optional.empty()); - expect(tierManager.getTier(TASK.getAssignedTask().getTask())).andReturn(DEV_TIER); expectLastCall(); control.replay(); @@ -299,55 +286,48 @@ public class TaskAssignerImplTest extends EasyMockTest { assigner.maybeAssign( storeProvider, resourceRequest, - TaskGroupKey.from(TASK.getAssignedTask().getTask()), - ImmutableSet.of(TASK.getAssignedTask()), + TaskGroupKey.from(TASK.getTask()), + ImmutableSet.of(TASK), ImmutableMap.of(SLAVE_ID, GROUP_KEY))); assertEquals(empty(), aggregate); } @Test public void testAssignWithMixOfReservedAndNotReserved() throws Exception { - expect(tierManager.getTier(TASK.getAssignedTask().getTask())).andReturn(DEV_TIER); - - expect(updateAgentReserver.hasReservations(GROUP_KEY)).andReturn(true); expect(updateAgentReserver.getAgent(INSTANCE_KEY)).andReturn(Optional.of(SLAVE_ID)); updateAgentReserver.release(SLAVE_ID, INSTANCE_KEY); - expect(offerManager.getMatching(MESOS_OFFER.getAgentId(), resourceRequest, false)) + expect(offerManager.getMatching(MESOS_OFFER.getAgentId(), resourceRequest)) .andReturn(Optional.of(OFFER)); expectAssignTask(MESOS_OFFER); offerManager.launchTask(MESOS_OFFER.getId(), TASK_INFO); - expect(taskFactory.createFrom(TASK.getAssignedTask(), MESOS_OFFER, false)) + expect(taskFactory.createFrom(TASK, MESOS_OFFER, false)) .andReturn(TASK_INFO); - // Normal scheduling loop for the remaining task... - IScheduledTask secondTask = makeTask("another-task", JOB, 9999); + // Normal scheduling loop for the remaining task. + IAssignedTask secondTask = makeTask("another-task", JOB, 9999).getAssignedTask(); TaskInfo secondTaskInfo = TaskInfo.newBuilder() .setName("another-task") - .setTaskId(TaskID.newBuilder().setValue(Tasks.id(secondTask))) + .setTaskId(TaskID.newBuilder().setValue(secondTask.getTaskId())) .setAgentId(MESOS_OFFER.getAgentId()) .build(); expect(updateAgentReserver.getAgent(InstanceKeys.from(JOB, 9999))).andReturn(Optional.empty()); - ImmutableSet<HostOffer> matchingOffers = ImmutableSet.of(OFFER); - expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest, false)) - .andReturn(matchingOffers); - expect(updateAgentReserver.getReservations(OFFER.getOffer().getAgentId().getValue())) - .andReturn(ImmutableSet.of()); - expectAssignTask(MESOS_OFFER, secondTask); - offerManager.launchTask(MESOS_OFFER.getId(), secondTaskInfo); - expect(taskFactory.createFrom(secondTask.getAssignedTask(), MESOS_OFFER, false)) - .andReturn(secondTaskInfo); + expect(offerManager.getAllMatching(GROUP_KEY, resourceRequest)) + .andReturn(ImmutableSet.of(OFFER_2)); + expect(updateAgentReserver.isReserved(OFFER_2.getOffer().getAgentId().getValue())) + .andReturn(false); + expectAssignTask(MESOS_OFFER_2, secondTask); + offerManager.launchTask(MESOS_OFFER_2.getId(), secondTaskInfo); + expect(taskFactory.createFrom(secondTask, MESOS_OFFER_2, false)).andReturn(secondTaskInfo); control.replay(); assertEquals( - Tasks.ids(TASK, secondTask), + ImmutableSet.of(TASK.getTaskId(), secondTask.getTaskId()), assigner.maybeAssign( storeProvider, resourceRequest, GROUP_KEY, - ImmutableSet.of( - TASK.getAssignedTask(), - secondTask.getAssignedTask()), + ImmutableSet.of(TASK, secondTask), ImmutableMap.of(SLAVE_ID, GROUP_KEY))); assertNotEquals(empty(), aggregate); } @@ -356,19 +336,19 @@ public class TaskAssignerImplTest extends EasyMockTest { expectAssignTask(offer, TASK); } - private void expectAssignTask(Offer offer, IScheduledTask task) { + private void expectAssignTask(Offer offer, IAssignedTask task) { expect(stateManager.assignTask( eq(storeProvider), - eq(Tasks.id(task)), + eq(task.getTaskId()), eq(offer.getHostname()), eq(offer.getAgentId()), - anyObject())).andReturn(task.getAssignedTask()); + anyObject())).andReturn(task); } private void expectNoUpdateReservations(int offers) { - expect(updateAgentReserver.hasReservations(anyObject())).andReturn(false); - for (int i = 0; i < offers; i++) { - expect(updateAgentReserver.getReservations(anyString())).andReturn(ImmutableSet.of()); + if (offers > 0) { + expect(updateAgentReserver.isReserved(anyString())).andReturn(false).times(offers); } + expect(updateAgentReserver.getAgent(anyObject())).andReturn(Optional.empty()); } } http://git-wip-us.apache.org/repos/asf/aurora/blob/4e6242fe/src/test/java/org/apache/aurora/scheduler/scheduling/TaskSchedulerImplTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/scheduling/TaskSchedulerImplTest.java b/src/test/java/org/apache/aurora/scheduler/scheduling/TaskSchedulerImplTest.java index ecf2987..ac2df94 100644 --- a/src/test/java/org/apache/aurora/scheduler/scheduling/TaskSchedulerImplTest.java +++ b/src/test/java/org/apache/aurora/scheduler/scheduling/TaskSchedulerImplTest.java @@ -31,6 +31,7 @@ import org.apache.aurora.common.stats.StatsProvider; import org.apache.aurora.common.testing.easymock.EasyMockTest; import org.apache.aurora.common.util.Clock; import org.apache.aurora.gen.ScheduledTask; +import org.apache.aurora.scheduler.TierManager; import org.apache.aurora.scheduler.async.AsyncModule.AsyncExecutor; import org.apache.aurora.scheduler.base.JobKeys; import org.apache.aurora.scheduler.base.Query; @@ -44,8 +45,6 @@ import org.apache.aurora.scheduler.events.PubsubEventModule; import org.apache.aurora.scheduler.filter.SchedulingFilter.ResourceRequest; import org.apache.aurora.scheduler.preemptor.BiCache; import org.apache.aurora.scheduler.preemptor.Preemptor; -import org.apache.aurora.scheduler.resources.ResourceBag; -import org.apache.aurora.scheduler.resources.ResourceManager; import org.apache.aurora.scheduler.state.PubsubTestUtil; import org.apache.aurora.scheduler.storage.Storage; import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult; @@ -62,6 +61,7 @@ import org.junit.Test; import static org.apache.aurora.gen.ScheduleStatus.PENDING; import static org.apache.aurora.gen.ScheduleStatus.RUNNING; import static org.apache.aurora.gen.ScheduleStatus.THROTTLED; +import static org.apache.aurora.scheduler.base.TaskTestUtil.TIER_MANAGER; import static org.apache.aurora.scheduler.filter.AttributeAggregate.empty; import static org.apache.aurora.scheduler.mesos.TestExecutorSettings.THERMOS_EXECUTOR; import static org.easymock.EasyMock.eq; @@ -105,7 +105,7 @@ public class TaskSchedulerImplTest extends EasyMockTest { new AbstractModule() { @Override protected void configure() { - + bind(TierManager.class).toInstance(TIER_MANAGER); bind(Executor.class).annotatedWith(AsyncExecutor.class) .toInstance(MoreExecutors.directExecutor()); bind(new TypeLiteral<BiCache<String, TaskGroupKey>>() { }).toInstance(reservations); @@ -127,28 +127,24 @@ public class TaskSchedulerImplTest extends EasyMockTest { ImmutableSet.of(task)); } - private ResourceBag bag(IScheduledTask task) { - return ResourceManager.bagFromResources(task.getAssignedTask().getTask().getResources()) - .add(THERMOS_EXECUTOR.getExecutorOverhead(task.getAssignedTask() - .getTask() - .getExecutorConfig() - .getName()).get()); - } - private IExpectationSetters<Set<String>> expectAssigned( IScheduledTask task, Map<String, TaskGroupKey> reservationMap) { return expect(assigner.maybeAssign( storageUtil.mutableStoreProvider, - new ResourceRequest(task.getAssignedTask().getTask(), bag(task), empty()), + ResourceRequest.fromTask( + task.getAssignedTask().getTask(), + THERMOS_EXECUTOR, + empty(), + TIER_MANAGER), TaskGroupKey.from(task.getAssignedTask().getTask()), ImmutableSet.of(task.getAssignedTask()), reservationMap)); } @Test - public void testSchedule() throws Exception { + public void testSchedule() { storageUtil.expectOperations(); expectAsMap(NO_RESERVATION); @@ -164,7 +160,7 @@ public class TaskSchedulerImplTest extends EasyMockTest { } @Test - public void testScheduleNoTask() throws Exception { + public void testScheduleNoTask() { storageUtil.expectOperations(); storageUtil.expectTaskFetch( Query.taskScoped(Tasks.id(TASK_A)).byStatus(PENDING), @@ -178,7 +174,7 @@ public class TaskSchedulerImplTest extends EasyMockTest { } @Test - public void testSchedulePartial() throws Exception { + public void testSchedulePartial() { storageUtil.expectOperations(); String taskB = "b"; @@ -214,7 +210,7 @@ public class TaskSchedulerImplTest extends EasyMockTest { } @Test - public void testReservation() throws Exception { + public void testReservation() { storageUtil.expectOperations(); // No reservation available in preemptor @@ -254,7 +250,7 @@ public class TaskSchedulerImplTest extends EasyMockTest { } @Test - public void testReservationUnusable() throws Exception { + public void testReservationUnusable() { storageUtil.expectOperations(); expectTaskStillPendingQuery(TASK_A); @@ -271,7 +267,7 @@ public class TaskSchedulerImplTest extends EasyMockTest { } @Test - public void testReservationRemoved() throws Exception { + public void testReservationRemoved() { storageUtil.expectOperations(); expectTaskStillPendingQuery(TASK_A); @@ -288,14 +284,14 @@ public class TaskSchedulerImplTest extends EasyMockTest { } @Test - public void testNonPendingIgnored() throws Exception { + public void testNonPendingIgnored() { control.replay(); eventSink.post(TaskStateChange.transition(TASK_A, RUNNING)); } @Test - public void testPendingDeletedHandled() throws Exception { + public void testPendingDeletedHandled() { reservations.remove(SLAVE_ID, TaskGroupKey.from(TASK_A.getAssignedTask().getTask())); control.replay(); @@ -325,7 +321,11 @@ public class TaskSchedulerImplTest extends EasyMockTest { expectAsMap(NO_RESERVATION); expect(assigner.maybeAssign( EasyMock.anyObject(), - eq(new ResourceRequest(taskA.getAssignedTask().getTask(), bag(taskA), empty())), + eq(ResourceRequest.fromTask( + taskA.getAssignedTask().getTask(), + THERMOS_EXECUTOR, + empty(), + TIER_MANAGER)), eq(TaskGroupKey.from(taskA.getAssignedTask().getTask())), eq(ImmutableSet.of(taskA.getAssignedTask())), eq(NO_RESERVATION))).andReturn(SCHEDULED_RESULT); @@ -337,7 +337,7 @@ public class TaskSchedulerImplTest extends EasyMockTest { } @Test - public void testScheduleThrows() throws Exception { + public void testScheduleThrows() { storageUtil.expectOperations(); expectAsMap(NO_RESERVATION); @@ -370,17 +370,17 @@ public class TaskSchedulerImplTest extends EasyMockTest { reservations.put(slaveId, TaskGroupKey.from(task.getAssignedTask().getTask())); } - private IExpectationSetters<?> expectGetReservation(IScheduledTask task, String slaveId) { - return expect(reservations.getByValue(TaskGroupKey.from(task.getAssignedTask().getTask()))) + private void expectGetReservation(IScheduledTask task, String slaveId) { + expect(reservations.getByValue(TaskGroupKey.from(task.getAssignedTask().getTask()))) .andReturn(ImmutableSet.of(slaveId)); } - private IExpectationSetters<?> expectNoReservation(IScheduledTask task) { - return expect(reservations.getByValue(TaskGroupKey.from(task.getAssignedTask().getTask()))) + private void expectNoReservation(IScheduledTask task) { + expect(reservations.getByValue(TaskGroupKey.from(task.getAssignedTask().getTask()))) .andReturn(ImmutableSet.of()); } - private IExpectationSetters<?> expectAsMap(Map<String, TaskGroupKey> map) { - return expect(reservations.asMap()).andReturn(map); + private void expectAsMap(Map<String, TaskGroupKey> map) { + expect(reservations.asMap()).andReturn(map); } } http://git-wip-us.apache.org/repos/asf/aurora/blob/4e6242fe/src/test/java/org/apache/aurora/scheduler/updater/NullAgentReserverTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/updater/NullAgentReserverTest.java b/src/test/java/org/apache/aurora/scheduler/updater/NullAgentReserverTest.java index d8563b8..6eebdb5 100644 --- a/src/test/java/org/apache/aurora/scheduler/updater/NullAgentReserverTest.java +++ b/src/test/java/org/apache/aurora/scheduler/updater/NullAgentReserverTest.java @@ -21,7 +21,6 @@ import org.apache.aurora.scheduler.updater.UpdateAgentReserver.NullAgentReserver import org.junit.Test; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertTrue; public class NullAgentReserverTest extends EasyMockTest { private static final IInstanceKey INSTANCE_KEY = @@ -33,6 +32,6 @@ public class NullAgentReserverTest extends EasyMockTest { NullAgentReserver reserver = new NullAgentReserver(); reserver.reserve("test", INSTANCE_KEY); assertFalse(reserver.getAgent(INSTANCE_KEY).isPresent()); - assertTrue(reserver.getReservations("test").isEmpty()); + assertFalse(reserver.isReserved("test")); } } http://git-wip-us.apache.org/repos/asf/aurora/blob/4e6242fe/src/test/java/org/apache/aurora/scheduler/updater/UpdateAgentReserverImplTest.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/aurora/scheduler/updater/UpdateAgentReserverImplTest.java b/src/test/java/org/apache/aurora/scheduler/updater/UpdateAgentReserverImplTest.java index 7f17be0..051ac0e 100644 --- a/src/test/java/org/apache/aurora/scheduler/updater/UpdateAgentReserverImplTest.java +++ b/src/test/java/org/apache/aurora/scheduler/updater/UpdateAgentReserverImplTest.java @@ -13,25 +13,19 @@ */ package org.apache.aurora.scheduler.updater; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import org.apache.aurora.common.testing.easymock.EasyMockTest; -import org.apache.aurora.gen.Resource; -import org.apache.aurora.gen.TaskConfig; import org.apache.aurora.scheduler.base.InstanceKeys; import org.apache.aurora.scheduler.base.JobKeys; -import org.apache.aurora.scheduler.base.TaskGroupKey; import org.apache.aurora.scheduler.preemptor.BiCache; import org.apache.aurora.scheduler.storage.entities.IInstanceKey; -import org.apache.aurora.scheduler.storage.entities.ITaskConfig; import org.apache.aurora.scheduler.updater.UpdateAgentReserver.UpdateAgentReserverImpl; import org.junit.Before; import org.junit.Test; import static org.easymock.EasyMock.expect; import static org.easymock.EasyMock.expectLastCall; -import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -43,16 +37,6 @@ public class UpdateAgentReserverImplTest extends EasyMockTest { private static final IInstanceKey INSTANCE_KEY = InstanceKeys.from(JobKeys.from("role", "env", "name"), 1); - private TaskGroupKey getTaskGroup(IInstanceKey key) { - return TaskGroupKey.from(ITaskConfig.build( - new TaskConfig() - .setJob(key.getJobKey().newBuilder()) - .setResources(ImmutableSet.of( - Resource.numCpus(1.0), - Resource.ramMb(1L), - Resource.diskMb(1L))))); - } - @Before public void setUp() { cache = createMock(new Clazz<BiCache<IInstanceKey, String>>() { }); @@ -76,32 +60,11 @@ public class UpdateAgentReserverImplTest extends EasyMockTest { } @Test - public void testGetReservations() { + public void testIsReserved() { expect(cache.getByValue(AGENT_ID)).andReturn(ImmutableSet.of(INSTANCE_KEY)); + expect(cache.getByValue(AGENT_ID)).andReturn(ImmutableSet.of()); control.replay(); - assertEquals(ImmutableSet.of(INSTANCE_KEY), reserver.getReservations(AGENT_ID)); + assertTrue(reserver.isReserved(AGENT_ID)); + assertFalse(reserver.isReserved(AGENT_ID)); } - - @Test - public void testHasReservations() { - IInstanceKey instanceKey2 = InstanceKeys.from(JobKeys.from("role", "env", "name"), 2); - IInstanceKey instanceKey3 = InstanceKeys.from(JobKeys.from("role2", "env2", "name2"), 1); - expect(cache.asMap()) - .andReturn(ImmutableMap.of( - INSTANCE_KEY, - AGENT_ID, - instanceKey2, - AGENT_ID, - instanceKey3, - "different-agent")).anyTimes(); - control.replay(); - assertTrue(reserver.hasReservations(getTaskGroup(INSTANCE_KEY))); - assertTrue(reserver.hasReservations(getTaskGroup(instanceKey2))); - assertTrue(reserver.hasReservations(getTaskGroup(instanceKey3))); - assertTrue(reserver.hasReservations( - getTaskGroup(InstanceKeys.from(JobKeys.from("role", "env", "name"), 3)))); - assertFalse(reserver.hasReservations( - getTaskGroup(InstanceKeys.from(JobKeys.from("not", "in", "map"), 1)))); - } - }
