aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From wfar...@apache.org
Subject [1/2] aurora git commit: Refactor scheduling code to split matching and assigning phases
Date Tue, 09 Jan 2018 22:51:09 GMT
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))));
-  }
-
 }


Mime
View raw message