aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From wfar...@apache.org
Subject [3/3] aurora git commit: Require non-default primitive values in StorageEntityUtil.
Date Tue, 21 Apr 2015 18:54:52 GMT
Require non-default primitive values in StorageEntityUtil.

Reviewed at https://reviews.apache.org/r/33200/


Project: http://git-wip-us.apache.org/repos/asf/aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/449a835e
Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/449a835e
Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/449a835e

Branch: refs/heads/master
Commit: 449a835e9a446987cd9033345eb580030ed35cfa
Parents: 3380572
Author: Bill Farner <wfarner@apache.org>
Authored: Tue Apr 21 11:51:51 2015 -0700
Committer: Bill Farner <wfarner@apache.org>
Committed: Tue Apr 21 11:51:51 2015 -0700

----------------------------------------------------------------------
 .../java/org/apache/aurora/benchmark/Tasks.java |   39 +-
 .../aurora/scheduler/base/TaskTestUtil.java     |  113 ++
 .../scheduler/async/TaskSchedulerImplTest.java  |   25 +-
 .../scheduler/async/TaskSchedulerTest.java      |   93 +-
 .../apache/aurora/scheduler/base/JobsTest.java  |    5 +-
 .../aurora/scheduler/base/TaskTestUtil.java     |   66 --
 .../apache/aurora/scheduler/base/TasksTest.java |   21 +-
 .../scheduler/state/StateManagerImplTest.java   |  135 +--
 .../scheduler/stats/ResourceCounterTest.java    |   21 +-
 .../scheduler/storage/StorageBackfillTest.java  |  140 +--
 .../scheduler/storage/backup/RecoveryTest.java  |   25 +-
 .../storage/db/DBJobUpdateStoreTest.java        | 1076 ------------------
 .../storage/db/DbJobUpdateStoreTest.java        | 1061 +++++++++++++++++
 .../scheduler/storage/mem/MemTaskStoreTest.java |  391 ++++---
 .../storage/testing/StorageEntityUtil.java      |   27 +-
 .../scheduler/updater/InstanceUpdaterTest.java  |   35 +
 .../aurora/scheduler/updater/JobUpdaterIT.java  |   25 +-
 .../aurora/scheduler/updater/TaskUtil.java      |   64 --
 18 files changed, 1665 insertions(+), 1697 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/jmh/java/org/apache/aurora/benchmark/Tasks.java
----------------------------------------------------------------------
diff --git a/src/jmh/java/org/apache/aurora/benchmark/Tasks.java b/src/jmh/java/org/apache/aurora/benchmark/Tasks.java
index 23827ab..b4da057 100644
--- a/src/jmh/java/org/apache/aurora/benchmark/Tasks.java
+++ b/src/jmh/java/org/apache/aurora/benchmark/Tasks.java
@@ -20,17 +20,16 @@ import com.google.common.collect.Lists;
 import com.twitter.common.quantity.Amount;
 import com.twitter.common.quantity.Data;
 
-import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.Constraint;
-import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.LimitConstraint;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.gen.TaskConstraint;
 import org.apache.aurora.gen.TaskEvent;
 import org.apache.aurora.gen.ValueConstraint;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
+import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 
 /**
@@ -134,25 +133,23 @@ final class Tasks {
         String taskId =
             jobKey.getRole() + "-" + jobKey.getEnvironment() + "-" + i + "-" + (uuidStart + i);
 
-        tasks.add(IScheduledTask.build(new ScheduledTask()
-            .setTaskEvents(Lists.newArrayList(new TaskEvent(0, ScheduleStatus.PENDING)))
+        ScheduledTask builder = TaskTestUtil.makeTask(taskId, IJobKey.build(jobKey))
+            .newBuilder()
             .setStatus(scheduleStatus)
-            .setAssignedTask(new AssignedTask()
-                .setInstanceId(i)
-                .setTaskId(taskId)
-                .setTask(new TaskConfig()
-                    .setConstraints(constraints.build())
-                    .setNumCpus(cpu)
-                    .setRamMb(ram.as(Data.MB))
-                    .setDiskMb(disk.as(Data.MB))
-                    .setProduction(isProduction)
-                    .setRequestedPorts(ImmutableSet.<String>of())
-                    .setJob(jobKey)
-                    .setJobName(jobKey.getName())
-                    .setEnvironment(jobKey.getEnvironment())
-                    .setOwner(new Identity()
-                        .setRole(jobKey.getRole())
-                        .setUser(String.format(USER_FORMAT, taskId)))))));
+            .setTaskEvents(Lists.newArrayList(
+                new TaskEvent(0, ScheduleStatus.PENDING),
+                new TaskEvent(1, scheduleStatus)));
+        builder.getAssignedTask()
+            .setInstanceId(i)
+            .setTaskId(taskId);
+        builder.getAssignedTask().getTask()
+            .setConstraints(constraints.build())
+            .setNumCpus(cpu)
+            .setRamMb(ram.as(Data.MB))
+            .setDiskMb(disk.as(Data.MB))
+            .setProduction(isProduction)
+            .setRequestedPorts(ImmutableSet.<String>of());
+        tasks.add(IScheduledTask.build(builder));
       }
 
       return tasks.build();

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java b/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
new file mode 100644
index 0000000..2f42851
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
@@ -0,0 +1,113 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.aurora.scheduler.base;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+
+import org.apache.aurora.gen.AssignedTask;
+import org.apache.aurora.gen.Constraint;
+import org.apache.aurora.gen.Container;
+import org.apache.aurora.gen.DockerContainer;
+import org.apache.aurora.gen.ExecutorConfig;
+import org.apache.aurora.gen.Identity;
+import org.apache.aurora.gen.LimitConstraint;
+import org.apache.aurora.gen.Metadata;
+import org.apache.aurora.gen.ScheduleStatus;
+import org.apache.aurora.gen.ScheduledTask;
+import org.apache.aurora.gen.TaskConfig;
+import org.apache.aurora.gen.TaskConstraint;
+import org.apache.aurora.gen.TaskEvent;
+import org.apache.aurora.gen.ValueConstraint;
+import org.apache.aurora.scheduler.storage.entities.IJobKey;
+import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
+import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
+
+/**
+ * Convenience methods for working with tasks.
+ * <p>
+ * TODO(wfarner): This lives in under src/main only so benchmarks can access it.  Reconsider the
+ * project layout so this is not necessary.
+ */
+public final class TaskTestUtil {
+
+  public static final IJobKey JOB = JobKeys.from("role", "env", "job");
+
+  private TaskTestUtil() {
+    // Utility class.
+  }
+
+  public static ITaskConfig makeConfig(IJobKey job) {
+    return ITaskConfig.build(new TaskConfig()
+        .setJob(job.newBuilder())
+        .setJobName(job.getName())
+        .setEnvironment(job.getEnvironment())
+        .setOwner(new Identity(job.getRole(), job.getRole() + "-user"))
+        .setIsService(true)
+        .setNumCpus(1.0)
+        .setRamMb(1024)
+        .setDiskMb(1024)
+        .setPriority(1)
+        .setMaxTaskFailures(-1)
+        .setProduction(true)
+        .setConstraints(ImmutableSet.of(
+            new Constraint(
+                "valueConstraint",
+                TaskConstraint.value(
+                    new ValueConstraint(false, ImmutableSet.of("value1", "value2")))),
+            new Constraint(
+                "limitConstraint",
+                TaskConstraint.limit(new LimitConstraint(5)))))
+        .setRequestedPorts(ImmutableSet.of("http"))
+        .setTaskLinks(ImmutableMap.of("http", "link", "admin", "otherLink"))
+        .setContactEmail("foo@bar.com")
+        .setMetadata(ImmutableSet.of(new Metadata("key", "value")))
+        .setExecutorConfig(new ExecutorConfig("name", "config"))
+        .setContainer(Container.docker(new DockerContainer("imagename"))));
+  }
+
+  public static IScheduledTask makeTask(String id, IJobKey job) {
+    return IScheduledTask.build(new ScheduledTask()
+        .setStatus(ScheduleStatus.PENDING)
+        .setTaskEvents(ImmutableList.of(
+            new TaskEvent(100L, ScheduleStatus.PENDING)
+                .setMessage("message")
+                .setScheduler("scheduler"),
+            new TaskEvent(101L, ScheduleStatus.ASSIGNED)
+                .setMessage("message")
+                .setScheduler("scheduler2")))
+        .setAncestorId("ancestor")
+        .setFailureCount(3)
+        .setAssignedTask(new AssignedTask()
+            .setInstanceId(2)
+            .setTaskId(id)
+            .setAssignedPorts(ImmutableMap.of("http", 1000))
+            .setTask(makeConfig(job).newBuilder())));
+  }
+
+  public static IScheduledTask addStateTransition(
+      IScheduledTask task,
+      ScheduleStatus status,
+      long timestamp) {
+
+    ScheduledTask builder = task.newBuilder();
+    builder.setStatus(status);
+    builder.addToTaskEvents(new TaskEvent()
+        .setTimestamp(timestamp)
+        .setStatus(status)
+        .setScheduler("scheduler"));
+    return IScheduledTask.build(builder);
+  }
+}

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerImplTest.java b/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerImplTest.java
index d879d33..bd1aaa1 100644
--- a/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerImplTest.java
@@ -20,24 +20,21 @@ import com.google.inject.AbstractModule;
 import com.google.inject.Guice;
 import com.google.inject.Injector;
 import com.google.inject.TypeLiteral;
-
 import com.twitter.common.stats.StatsProvider;
 import com.twitter.common.testing.easymock.EasyMockTest;
 import com.twitter.common.util.Clock;
 
-import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.HostAttributes;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.scheduler.HostOffer;
 import org.apache.aurora.scheduler.async.TaskScheduler.TaskSchedulerImpl;
 import org.apache.aurora.scheduler.async.preemptor.BiCache;
 import org.apache.aurora.scheduler.async.preemptor.Preemptor;
+import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.TaskGroupKey;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.events.EventSink;
 import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
@@ -75,8 +72,10 @@ import static org.junit.Assert.assertTrue;
 
 public class TaskSchedulerImplTest extends EasyMockTest {
 
-  private static final IScheduledTask TASK_A = makeTask("a");
-  private static final IScheduledTask TASK_B = makeTask("b");
+  private static final IScheduledTask TASK_A =
+      TaskTestUtil.makeTask("a", JobKeys.from("a", "a", "a"));
+  private static final IScheduledTask TASK_B =
+      TaskTestUtil.makeTask("b", JobKeys.from("b", "b", "b"));
   private static final HostOffer OFFER = new HostOffer(
       Offers.makeOffer("OFFER_A", "HOST_A"),
       IHostAttributes.build(new HostAttributes().setMode(MaintenanceMode.NONE)));
@@ -276,18 +275,6 @@ public class TaskSchedulerImplTest extends EasyMockTest {
     assignAndAssert(Result.SUCCESS, GROUP_A, OFFER, assignment);
   }
 
-  private static IScheduledTask makeTask(String taskId) {
-    return IScheduledTask.build(new ScheduledTask()
-        .setAssignedTask(new AssignedTask()
-            .setInstanceId(0)
-            .setTaskId(taskId)
-            .setTask(new TaskConfig()
-                .setJob(new JobKey("role-" + taskId, "env-" + taskId, "job-" + taskId))
-                .setJobName("job-" + taskId)
-                .setOwner(new Identity().setRole("role-" + taskId).setUser("user-" + taskId))
-                .setEnvironment("env-" + taskId))));
-  }
-
   private static class AssignmentCapture {
     public Capture<Function<HostOffer, Assignment>> assigner = createCapture();
     public Capture<TaskGroupKey> groupKey = createCapture();

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerTest.java b/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerTest.java
index 798ff6e..9035484 100644
--- a/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/async/TaskSchedulerTest.java
@@ -17,6 +17,8 @@ import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 
+import javax.annotation.Nullable;
+
 import com.google.common.base.Function;
 import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
@@ -28,11 +30,8 @@ import com.twitter.common.quantity.Time;
 import com.twitter.common.testing.easymock.EasyMockTest;
 import com.twitter.common.util.BackoffStrategy;
 
-import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.Attribute;
 import org.apache.aurora.gen.HostAttributes;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
@@ -43,8 +42,10 @@ import org.apache.aurora.scheduler.async.OfferManager.OfferReturnDelay;
 import org.apache.aurora.scheduler.async.TaskScheduler.TaskSchedulerImpl;
 import org.apache.aurora.scheduler.async.preemptor.BiCache;
 import org.apache.aurora.scheduler.async.preemptor.Preemptor;
+import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.TaskGroupKey;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.events.PubsubEvent;
 import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
@@ -61,6 +62,7 @@ import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
 import org.apache.aurora.scheduler.storage.Storage.MutateWork;
 import org.apache.aurora.scheduler.storage.Storage.StorageException;
 import org.apache.aurora.scheduler.storage.TaskStore;
+import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
@@ -85,6 +87,7 @@ import static org.apache.aurora.gen.ScheduleStatus.KILLED;
 import static org.apache.aurora.gen.ScheduleStatus.LOST;
 import static org.apache.aurora.gen.ScheduleStatus.PENDING;
 import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
+import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask;
 import static org.apache.aurora.scheduler.filter.AttributeAggregate.EMPTY;
 import static org.apache.mesos.Protos.Offer;
 import static org.easymock.EasyMock.capture;
@@ -221,37 +224,33 @@ public class TaskSchedulerTest extends EasyMockTest {
   public void testNoOffers() {
     Capture<Runnable> timeoutCapture = expectTaskRetryIn(FIRST_SCHEDULE_DELAY_MS);
     expectTaskGroupBackoff(FIRST_SCHEDULE_DELAY_MS, 10);
-    IScheduledTask task = makeTask("a");
-    expectPreemptorCall(task);
+    IScheduledTask task = createTask("a");
+    expectPreemptorCall(task.getAssignedTask());
     expectReservationCheck(task);
 
     replayAndCreateScheduler();
 
-    changeState(makeTask("a"), INIT, PENDING);
+    changeState(task, INIT, PENDING);
     timeoutCapture.getValue().run();
   }
 
-  private IScheduledTask makeTask(String taskId) {
-    return IScheduledTask.build(new ScheduledTask()
-        .setAssignedTask(new AssignedTask()
-            .setInstanceId(0)
-            .setTaskId(taskId)
-            .setTask(new TaskConfig()
-                .setJob(new JobKey("role-" + taskId, "env-" + taskId, "job-" + taskId))
-                .setJobName("job-" + taskId)
-                .setOwner(new Identity().setRole("role-" + taskId).setUser("user-" + taskId))
-                .setEnvironment("env-" + taskId))));
+  private IScheduledTask createTask(String taskId) {
+    return createTask(taskId, null);
+  }
+
+  private IScheduledTask createTask(String taskId, @Nullable ScheduleStatus status) {
+    return setStatus(makeTask(taskId, TaskTestUtil.JOB), status);
   }
 
-  private IScheduledTask makeTask(String taskId, ScheduleStatus status) {
-    return IScheduledTask.build(makeTask(taskId).newBuilder().setStatus(status));
+  private IScheduledTask setStatus(IScheduledTask task, @Nullable ScheduleStatus status) {
+    return IScheduledTask.build(task.newBuilder().setStatus(status));
   }
 
   @Test
   public void testLoadFromStorage() {
-    final IScheduledTask a = makeTask("a", KILLED);
-    final IScheduledTask b = makeTask("b", PENDING);
-    final IScheduledTask c = makeTask("c", RUNNING);
+    final IScheduledTask a = createTask("a", KILLED);
+    final IScheduledTask b = createTask("b", PENDING);
+    final IScheduledTask c = createTask("c", RUNNING);
 
     expect(rescheduleCalculator.getStartupScheduleDelayMs(b)).andReturn(10L);
     expectTaskRetryIn(10);
@@ -276,7 +275,7 @@ public class TaskSchedulerTest extends EasyMockTest {
 
     replayAndCreateScheduler();
 
-    taskGroups.taskChangedState(TaskStateChange.transition(makeTask("a", PENDING), INIT));
+    taskGroups.taskChangedState(TaskStateChange.transition(createTask("a", PENDING), INIT));
     timeoutCapture.getValue().run();
   }
 
@@ -306,40 +305,40 @@ public class TaskSchedulerTest extends EasyMockTest {
     expectAnyMaintenanceCalls();
     expectOfferDeclineIn(10);
 
-    IScheduledTask task = makeTask("a", PENDING);
-    TaskInfo mesosTask = makeTaskInfo(task);
+    IScheduledTask taskA = createTask("a", PENDING);
+    TaskInfo mesosTask = makeTaskInfo(taskA);
 
     Capture<Runnable> timeoutCapture = expectTaskRetryIn(FIRST_SCHEDULE_DELAY_MS);
     expectNoReservation(SLAVE_A).times(2);
-    expectReservationCheck(task);
-    expectMaybeAssign(OFFER_A, task, EMPTY).andReturn(Assignment.failure());
-    expectPreemptorCall(task);
+    expectReservationCheck(taskA);
+    expectMaybeAssign(OFFER_A, taskA, EMPTY).andReturn(Assignment.failure());
+    expectPreemptorCall(taskA.getAssignedTask());
 
     Capture<Runnable> timeoutCapture2 = expectTaskGroupBackoff(FIRST_SCHEDULE_DELAY_MS, 10);
-    expectMaybeAssign(OFFER_A, task, EMPTY).andReturn(Assignment.success(mesosTask));
+    expectMaybeAssign(OFFER_A, taskA, EMPTY).andReturn(Assignment.success(mesosTask));
     driver.launchTask(OFFER_A.getOffer().getId(), mesosTask);
 
     Capture<Runnable> timeoutCapture3 = expectTaskRetryIn(FIRST_SCHEDULE_DELAY_MS);
     expectTaskGroupBackoff(FIRST_SCHEDULE_DELAY_MS, 10);
-    IScheduledTask taskB = makeTask("b");
+    IScheduledTask taskB = createTask("b");
     expectReservationCheck(taskB);
-    expectPreemptorCall(taskB);
+    expectPreemptorCall(taskB.getAssignedTask());
 
     replayAndCreateScheduler();
 
     offerManager.addOffer(OFFER_A);
-    changeState(task, INIT, PENDING);
+    changeState(taskA, INIT, PENDING);
     timeoutCapture.getValue().run();
     timeoutCapture2.getValue().run();
 
     // Ensure the offer was consumed.
-    changeState(makeTask("b"), INIT, PENDING);
+    changeState(taskB, INIT, PENDING);
     timeoutCapture3.getValue().run();
   }
 
   @Test
   public void testDriverNotReady() {
-    IScheduledTask task = makeTask("a", PENDING);
+    IScheduledTask task = createTask("a", PENDING);
     TaskInfo mesosTask = TaskInfo.newBuilder()
         .setName(Tasks.id(task))
         .setTaskId(TaskID.newBuilder().setValue(Tasks.id(task)))
@@ -370,7 +369,7 @@ public class TaskSchedulerTest extends EasyMockTest {
 
   @Test
   public void testStorageException() {
-    IScheduledTask task = makeTask("a", PENDING);
+    IScheduledTask task = createTask("a", PENDING);
     TaskInfo mesosTask = TaskInfo.newBuilder()
         .setName(Tasks.id(task))
         .setTaskId(TaskID.newBuilder().setValue(Tasks.id(task)))
@@ -398,7 +397,7 @@ public class TaskSchedulerTest extends EasyMockTest {
 
   @Test
   public void testExpiration() {
-    IScheduledTask task = makeTask("a", PENDING);
+    IScheduledTask task = createTask("a", PENDING);
 
     Capture<Runnable> timeoutCapture = expectTaskRetryIn(FIRST_SCHEDULE_DELAY_MS);
     Capture<Runnable> offerExpirationCapture = expectOfferDeclineIn(10);
@@ -407,10 +406,10 @@ public class TaskSchedulerTest extends EasyMockTest {
     expectReservationCheck(task).times(2);
     expectMaybeAssign(OFFER_A, task, EMPTY).andReturn(Assignment.failure());
     Capture<Runnable> timeoutCapture2 = expectTaskGroupBackoff(FIRST_SCHEDULE_DELAY_MS, 10);
-    expectPreemptorCall(task);
+    expectPreemptorCall(task.getAssignedTask());
     driver.declineOffer(OFFER_A.getOffer().getId());
     expectTaskGroupBackoff(10, 20);
-    expectPreemptorCall(task);
+    expectPreemptorCall(task.getAssignedTask());
 
     replayAndCreateScheduler();
 
@@ -465,14 +464,14 @@ public class TaskSchedulerTest extends EasyMockTest {
     expectOffer();
     expectOffer();
 
-    IScheduledTask taskA = makeTask("A", PENDING);
+    IScheduledTask taskA = createTask("A", PENDING);
     TaskInfo mesosTaskA = makeTaskInfo(taskA);
     expectNoReservation(SLAVE_A);
     expectMaybeAssign(OFFER_A, taskA, EMPTY).andReturn(Assignment.success(mesosTaskA));
     driver.launchTask(OFFER_A.getOffer().getId(), mesosTaskA);
     Capture<Runnable> captureA = expectTaskRetryIn(FIRST_SCHEDULE_DELAY_MS);
 
-    IScheduledTask taskB = makeTask("B", PENDING);
+    IScheduledTask taskB = createTask("B", PENDING);
     TaskInfo mesosTaskB = makeTaskInfo(taskB);
     expectNoReservation(SLAVE_B);
     expectMaybeAssign(OFFER_B, taskB, EMPTY).andReturn(Assignment.success(mesosTaskB));
@@ -499,14 +498,14 @@ public class TaskSchedulerTest extends EasyMockTest {
     expectOffer();
     expectOffer();
 
-    IScheduledTask taskA = makeTask("A", PENDING);
+    IScheduledTask taskA = createTask("A", PENDING);
     TaskInfo mesosTaskA = makeTaskInfo(taskA);
     expectNoReservation(SLAVE_B);
     expectMaybeAssign(OFFER_B, taskA, EMPTY).andReturn(Assignment.success(mesosTaskA));
     driver.launchTask(OFFER_B.getOffer().getId(), mesosTaskA);
     Capture<Runnable> captureA = expectTaskRetryIn(FIRST_SCHEDULE_DELAY_MS);
 
-    IScheduledTask taskB = makeTask("B", PENDING);
+    IScheduledTask taskB = createTask("B", PENDING);
     TaskInfo mesosTaskB = makeTaskInfo(taskB);
     HostOffer updatedOfferC = new HostOffer(
         OFFER_C.getOffer(),
@@ -554,7 +553,7 @@ public class TaskSchedulerTest extends EasyMockTest {
 
     expectAnyMaintenanceCalls();
 
-    IScheduledTask jobA0 = makeTask("a0", PENDING);
+    IScheduledTask jobA0 = setStatus(makeTask("a0", JobKeys.from("a", "b", "c")), PENDING);
 
     ScheduledTask jobA1Builder = jobA0.newBuilder();
     jobA1Builder.getAssignedTask().setTaskId("a1");
@@ -566,7 +565,7 @@ public class TaskSchedulerTest extends EasyMockTest {
     jobA2Builder.getAssignedTask().setInstanceId(2);
     IScheduledTask jobA2 = IScheduledTask.build(jobA2Builder);
 
-    IScheduledTask jobB0 = makeTask("b0", PENDING);
+    IScheduledTask jobB0 = setStatus(makeTask("b0", JobKeys.from("d", "e", "f")), PENDING);
 
     expectNoReservation(SLAVE_A);
     expectNoReservation(SLAVE_B);
@@ -607,14 +606,14 @@ public class TaskSchedulerTest extends EasyMockTest {
     expectAnyMaintenanceCalls();
     expectOfferDeclineIn(10);
 
-    final IScheduledTask task = makeTask("a", PENDING);
+    final IScheduledTask task = createTask("a", PENDING);
 
     Capture<Runnable> timeoutCapture = expectTaskRetryIn(FIRST_SCHEDULE_DELAY_MS);
     expectNoReservation(SLAVE_A);
     expectMaybeAssign(OFFER_A, task, EMPTY).andReturn(Assignment.failure());
     expectTaskGroupBackoff(FIRST_SCHEDULE_DELAY_MS, 20);
     expectReservationCheck(task);
-    expectPreemptorCall(task);
+    expectPreemptorCall(task.getAssignedTask());
 
     replayAndCreateScheduler();
 
@@ -662,9 +661,9 @@ public class TaskSchedulerTest extends EasyMockTest {
             .setMode(mode)));
   }
 
-  private void expectPreemptorCall(IScheduledTask task) {
+  private void expectPreemptorCall(IAssignedTask task) {
     expect(preemptor.attemptPreemptionFor(
-        eq(task.getAssignedTask()),
+        eq(task),
         eq(EMPTY),
         EasyMock.<MutableStoreProvider>anyObject())).andReturn(Optional.<String>absent());
   }

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/base/JobsTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/base/JobsTest.java b/src/test/java/org/apache/aurora/scheduler/base/JobsTest.java
index a170fb6..986fd01 100644
--- a/src/test/java/org/apache/aurora/scheduler/base/JobsTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/base/JobsTest.java
@@ -25,12 +25,11 @@ import org.apache.aurora.scheduler.storage.entities.IJobStats;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.junit.Test;
 
+import static org.apache.aurora.scheduler.base.TaskTestUtil.addStateTransition;
 import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask;
-import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTaskEvents;
 import static org.junit.Assert.assertEquals;
 
 public class JobsTest {
-
   @Test
   public void testGetJobStats() {
     ImmutableList<IScheduledTask> tasks =
@@ -39,7 +38,7 @@ public class JobsTest {
             .transform(new Function<ScheduleStatus, IScheduledTask>() {
               @Override
               public IScheduledTask apply(ScheduleStatus status) {
-                return makeTask(status, makeTaskEvents(100, status));
+                return addStateTransition(makeTask("id", TaskTestUtil.JOB), status, 100L);
               }
             }).toList();
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/base/TaskTestUtil.java b/src/test/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
deleted file mode 100644
index b8b9c1b..0000000
--- a/src/test/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
+++ /dev/null
@@ -1,66 +0,0 @@
-/**
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.aurora.scheduler.base;
-
-import java.util.List;
-
-import com.google.common.collect.ImmutableList;
-
-import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.ScheduleStatus;
-import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
-import org.apache.aurora.gen.TaskEvent;
-import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
-
-/**
- * Convenience methods for working with tasks.
- */
-public final class TaskTestUtil {
-
-  private TaskTestUtil() {
-    // Utility class.
-  }
-
-  public static IScheduledTask makeTask(ScheduleStatus status, List<TaskEvent> taskEvents) {
-    return IScheduledTask.build(new ScheduledTask()
-        .setStatus(status)
-        .setAssignedTask(new AssignedTask()
-            .setTaskId("task_id")
-            .setSlaveHost("host1")
-            .setTask(new TaskConfig()
-                .setJobName("job_a")
-                .setOwner(new Identity("role_a", "role_a" + "-user"))))
-        .setTaskEvents(taskEvents));
-  }
-
-  public static List<TaskEvent> makeTaskEvents(
-      long startTs,
-      ScheduleStatus status,
-      ScheduleStatus... statuses) {
-
-    ImmutableList.Builder<TaskEvent> taskEvents = ImmutableList.builder();
-    taskEvents.add(makeTaskEvent(startTs, status));
-    long increment = startTs;
-    for (ScheduleStatus st : statuses) {
-      taskEvents.add(makeTaskEvent(startTs + increment++, st));
-    }
-    return taskEvents.build();
-  }
-
-  private static TaskEvent makeTaskEvent(long ts, ScheduleStatus status) {
-    return new TaskEvent().setTimestamp(ts).setStatus(status);
-  }
-}

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/base/TasksTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/base/TasksTest.java b/src/test/java/org/apache/aurora/scheduler/base/TasksTest.java
index 2866b20..ed9a287 100644
--- a/src/test/java/org/apache/aurora/scheduler/base/TasksTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/base/TasksTest.java
@@ -22,8 +22,6 @@ import org.junit.Test;
 
 import static org.apache.aurora.gen.ScheduleStatus.FINISHED;
 import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
-import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask;
-import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTaskEvents;
 import static org.apache.aurora.scheduler.base.Tasks.getLatestActiveTask;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
@@ -43,12 +41,12 @@ public class TasksTest {
 
   @Test
   public void testLatestTransitionedTasks() {
-    IScheduledTask f1 = makeTask(FINISHED, makeTaskEvents(100, FINISHED));
-    IScheduledTask f2 = makeTask(FINISHED, makeTaskEvents(200, FINISHED));
-    IScheduledTask f3 = makeTask(FINISHED, makeTaskEvents(300, FINISHED));
-    IScheduledTask r1 = makeTask(RUNNING, makeTaskEvents(400, RUNNING));
-    IScheduledTask r2 = makeTask(RUNNING, makeTaskEvents(500, RUNNING));
-    IScheduledTask r3 = makeTask(RUNNING, makeTaskEvents(600, RUNNING));
+    IScheduledTask f1 = makeTask(FINISHED, 100);
+    IScheduledTask f2 = makeTask(FINISHED, 200);
+    IScheduledTask f3 = makeTask(FINISHED, 300);
+    IScheduledTask r1 = makeTask(RUNNING, 400);
+    IScheduledTask r2 = makeTask(RUNNING, 500);
+    IScheduledTask r3 = makeTask(RUNNING, 600);
 
     try {
       getLatestActiveTask(ImmutableList.<IScheduledTask>of());
@@ -78,4 +76,11 @@ public class TasksTest {
   private void assertLatestTask(IScheduledTask expectedLatest, IScheduledTask... tasks) {
     assertEquals(expectedLatest, getLatestActiveTask(ImmutableList.copyOf(tasks)));
   }
+
+  private IScheduledTask makeTask(ScheduleStatus status, long timestamp) {
+    return TaskTestUtil.addStateTransition(
+        TaskTestUtil.makeTask("id-" + timestamp, TaskTestUtil.JOB),
+        status,
+        timestamp);
+  }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java b/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java
index 7234fd3..c7fd3e5 100644
--- a/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/state/StateManagerImplTest.java
@@ -29,22 +29,25 @@ import com.twitter.common.testing.easymock.EasyMockTest;
 import com.twitter.common.util.testing.FakeClock;
 
 import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
+import org.apache.aurora.gen.Attribute;
+import org.apache.aurora.gen.HostAttributes;
+import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.gen.TaskEvent;
 import org.apache.aurora.scheduler.TaskIdGenerator;
 import org.apache.aurora.scheduler.async.RescheduleCalculator;
 import org.apache.aurora.scheduler.base.Query;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.events.EventSink;
 import org.apache.aurora.scheduler.events.PubsubEvent;
 import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
 import org.apache.aurora.scheduler.events.PubsubEvent.TasksDeleted;
 import org.apache.aurora.scheduler.mesos.Driver;
+import org.apache.aurora.scheduler.storage.AttributeStore;
 import org.apache.aurora.scheduler.storage.Storage;
+import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
 import org.apache.aurora.scheduler.storage.mem.MemStorage;
@@ -75,9 +78,15 @@ import static org.junit.Assert.assertTrue;
 
 public class StateManagerImplTest extends EasyMockTest {
 
-  private static final String HOST_A = "host_a";
-  private static final Identity JIM = new Identity("jim", "jim-user");
-  private static final String MY_JOB = "myJob";
+  private static final IHostAttributes HOST_A = IHostAttributes.build(
+      new HostAttributes(
+          "hostA",
+          ImmutableSet.of(new Attribute("zone", ImmutableSet.of("1a"))))
+          .setSlaveId("slaveIdA")
+          .setMode(MaintenanceMode.NONE));
+  private static final ITaskConfig NON_SERVICE_CONFIG =
+      setIsService(TaskTestUtil.makeConfig(TaskTestUtil.JOB), false);
+  private static final ITaskConfig SERVICE_CONFIG = setIsService(NON_SERVICE_CONFIG, true);
 
   private Driver driver;
   private TaskIdGenerator taskIdGenerator;
@@ -101,6 +110,13 @@ public class StateManagerImplTest extends EasyMockTest {
         taskIdGenerator,
         eventSink,
         rescheduleCalculator);
+    storage.write(new Storage.MutateWork.NoResult.Quiet() {
+      @Override
+      protected void execute(Storage.MutableStoreProvider storeProvider) {
+        AttributeStore.Mutable attributeStore = storeProvider.getAttributeStore();
+        attributeStore.saveHostAttributes(HOST_A);
+      }
+    });
   }
 
   private static class StateChangeMatcher implements IArgumentMatcher {
@@ -167,14 +183,13 @@ public class StateManagerImplTest extends EasyMockTest {
 
   @Test
   public void testAddTasks() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 3)).andReturn(taskId);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 3)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING);
 
     control.replay();
 
-    insertTask(task, 3);
+    insertTask(NON_SERVICE_CONFIG, 3);
     ScheduledTask expected = new ScheduledTask()
         .setStatus(PENDING)
         .setTaskEvents(ImmutableList.of(new TaskEvent()
@@ -184,37 +199,36 @@ public class StateManagerImplTest extends EasyMockTest {
         .setAssignedTask(new AssignedTask()
             .setInstanceId(3)
             .setTaskId(taskId)
-            .setTask(task.newBuilder()));
-    assertEquals(ImmutableSet.of(IScheduledTask.build(expected)),
+            .setTask(NON_SERVICE_CONFIG.newBuilder()));
+    assertEquals(
+        ImmutableSet.of(IScheduledTask.build(expected)),
         Storage.Util.fetchTasks(storage, Query.taskScoped(taskId)));
   }
 
   @Test
   public void testKillPendingTask() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING);
     eventSink.post(matchTasksDeleted(taskId));
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
     assertEquals(true, changeState(taskId, KILLING));
     assertEquals(false, changeState(taskId, KILLING));
   }
 
   @Test
   public void testKillRunningTask() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING, ASSIGNED, RUNNING, KILLING, KILLED);
     driver.killTask(EasyMock.<String>anyObject());
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
     assignTask(taskId, HOST_A);
     assertEquals(true, changeState(taskId, RUNNING));
     assertEquals(true, changeState(taskId, KILLING));
@@ -224,16 +238,15 @@ public class StateManagerImplTest extends EasyMockTest {
 
   @Test
   public void testLostKillingTask() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING, ASSIGNED, RUNNING, KILLING, LOST);
 
     driver.killTask(EasyMock.<String>anyObject());
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
 
     assignTask(taskId, HOST_A);
     changeState(taskId, RUNNING);
@@ -244,8 +257,7 @@ public class StateManagerImplTest extends EasyMockTest {
   @Test
   public void testNestedEvents() {
     final String id = "a";
-    ITaskConfig task = makeTask(JIM, MY_JOB);
-    expect(taskIdGenerator.generate(task, 0)).andReturn(id);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(id);
 
     // Trigger an event that produces a side-effect and a PubSub event .
     eventSink.post(matchStateChange(id, INIT, PENDING));
@@ -262,26 +274,29 @@ public class StateManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
   }
 
   @Test
   public void testDeletePendingTask() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING);
     eventSink.post(matchTasksDeleted(taskId));
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
     changeState(taskId, KILLING);
   }
 
+  private static ITaskConfig setIsService(ITaskConfig config, boolean service) {
+    return ITaskConfig.build(config.newBuilder().setIsService(service));
+  }
+
   @Test
   public void testThrottleTask() {
-    ITaskConfig task = ITaskConfig.build(makeTask(JIM, MY_JOB).newBuilder().setIsService(true));
+    ITaskConfig task = setIsService(NON_SERVICE_CONFIG, true);
     String taskId = "a";
     expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING, ASSIGNED, RUNNING, FAILED);
@@ -317,19 +332,18 @@ public class StateManagerImplTest extends EasyMockTest {
 
   @Test
   public void testIncrementFailureCount() {
-    ITaskConfig task = ITaskConfig.build(makeTask(JIM, MY_JOB).newBuilder().setIsService(true));
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    expect(taskIdGenerator.generate(SERVICE_CONFIG, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING, ASSIGNED, RUNNING, FAILED);
 
     String taskId2 = "a2";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId2);
+    expect(taskIdGenerator.generate(SERVICE_CONFIG, 0)).andReturn(taskId2);
     noFlappingPenalty();
     expectStateTransitions(taskId2, INIT, PENDING);
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(SERVICE_CONFIG, 0);
 
     assignTask(taskId, HOST_A);
     changeState(taskId, RUNNING);
@@ -340,16 +354,20 @@ public class StateManagerImplTest extends EasyMockTest {
     assertEquals(1, rescheduledTask.getFailureCount());
   }
 
+  private static ITaskConfig setMaxFailures(ITaskConfig config, int maxFailures) {
+    return ITaskConfig.build(config.newBuilder().setMaxTaskFailures(maxFailures));
+  }
+
   @Test
   public void testCasTaskPresent() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    ITaskConfig config = setMaxFailures(NON_SERVICE_CONFIG, 1);
+    expect(taskIdGenerator.generate(config, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING, ASSIGNED, FAILED);
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(config, 0);
     assignTask(taskId, HOST_A);
     assertFalse(changeState(
         taskId,
@@ -376,15 +394,14 @@ public class StateManagerImplTest extends EasyMockTest {
 
   @Test
   public void testDeleteTasks() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     final String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId);
     expectStateTransitions(taskId, INIT, PENDING, ASSIGNED, RUNNING, FINISHED);
     eventSink.post(matchTasksDeleted(taskId));
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
     assignTask(taskId, HOST_A);
     changeState(taskId, RUNNING);
     changeState(taskId, FINISHED);
@@ -394,14 +411,16 @@ public class StateManagerImplTest extends EasyMockTest {
         stateManager.deleteTasks(storeProvider, ImmutableSet.of(taskId));
       }
     });
+  }
 
+  private static ITaskConfig setRequestedPorts(ITaskConfig config, Set<String> portNames) {
+    return ITaskConfig.build(config.newBuilder().setRequestedPorts(portNames));
   }
 
   @Test
   public void testPortResource() throws Exception {
     Set<String> requestedPorts = ImmutableSet.of("one", "two", "three");
-    ITaskConfig task = ITaskConfig.build(makeTask(JIM, MY_JOB).newBuilder()
-        .setRequestedPorts(requestedPorts));
+    ITaskConfig task = setRequestedPorts(NON_SERVICE_CONFIG, requestedPorts);
 
     String taskId = "a";
     expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
@@ -423,8 +442,7 @@ public class StateManagerImplTest extends EasyMockTest {
   @Test
   public void testPortResourceResetAfterReschedule() throws Exception {
     Set<String> requestedPorts = ImmutableSet.of("one");
-    ITaskConfig task = ITaskConfig.build(makeTask(JIM, MY_JOB).newBuilder()
-        .setRequestedPorts(requestedPorts));
+    ITaskConfig task = setRequestedPorts(NON_SERVICE_CONFIG, requestedPorts);
 
     String taskId = "a";
     expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
@@ -458,7 +476,7 @@ public class StateManagerImplTest extends EasyMockTest {
       protected void execute(Storage.MutableStoreProvider storeProvider) {
         stateManager.insertPendingTasks(
             storeProvider,
-            makeTask(JIM, MY_JOB),
+            NON_SERVICE_CONFIG,
             ImmutableSet.<Integer>of());
       }
     });
@@ -466,26 +484,24 @@ public class StateManagerImplTest extends EasyMockTest {
 
   @Test(expected = IllegalArgumentException.class)
   public void insertPendingInstancesInstanceCollision() {
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId).times(2);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId).times(2);
     expectStateTransitions(taskId, INIT, PENDING);
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
     Iterables.getOnlyElement(Storage.Util.fetchTasks(storage, Query.taskScoped(taskId)));
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
   }
 
   @Test
   public void testAssignTaskPubsub() {
     // This test ensures the pubsub events emitted by assigning tasks have slave id and host set.
-    ITaskConfig task = makeTask(JIM, MY_JOB);
     String taskId = "a";
 
-    expect(taskIdGenerator.generate(task, 0)).andReturn(taskId);
+    expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId);
 
     expectStateTransitions(taskId, INIT, PENDING);
 
@@ -494,13 +510,13 @@ public class StateManagerImplTest extends EasyMockTest {
 
     control.replay();
 
-    insertTask(task, 0);
+    insertTask(NON_SERVICE_CONFIG, 0);
     assignTask(taskId, HOST_A);
 
     TaskStateChange change = taskStateChangeCapture.getValue();
     assertEquals(ASSIGNED, change.getNewState());
-    assertEquals(HOST_A, change.getTask().getAssignedTask().getSlaveHost());
-    assertEquals(HOST_A, change.getTask().getAssignedTask().getSlaveId());
+    assertEquals(HOST_A.getHost(), change.getTask().getAssignedTask().getSlaveHost());
+    assertEquals(HOST_A.getSlaveId(), change.getTask().getAssignedTask().getSlaveId());
   }
 
   private void expectStateTransitions(
@@ -561,22 +577,13 @@ public class StateManagerImplTest extends EasyMockTest {
         Optional.<String>absent());
   }
 
-  private static ITaskConfig makeTask(Identity owner, String job) {
-    return ITaskConfig.build(new TaskConfig()
-        .setJob(new JobKey(owner.getRole(), "devel", job))
-        .setOwner(owner)
-        .setEnvironment("devel")
-        .setJobName(job)
-        .setRequestedPorts(ImmutableSet.<String>of()));
-  }
-
-  private void assignTask(String taskId, String host) {
+  private void assignTask(String taskId, IHostAttributes host) {
     assignTask(taskId, host, ImmutableMap.<String, Integer>of());
   }
 
   private void assignTask(
       final String taskId,
-      final String host,
+      final IHostAttributes host,
       final Map<String, Integer> ports) {
 
     storage.write(new Storage.MutateWork.NoResult.Quiet() {
@@ -585,8 +592,8 @@ public class StateManagerImplTest extends EasyMockTest {
         stateManager.assignTask(
             storeProvider,
             taskId,
-            host,
-            SlaveID.newBuilder().setValue(host).build(),
+            host.getHost(),
+            SlaveID.newBuilder().setValue(host.getSlaveId()).build(),
             ports);
       }
     });

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java b/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java
index e5f3adf..82e9c76 100644
--- a/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/stats/ResourceCounterTest.java
@@ -21,10 +21,7 @@ import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 
-import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.Constraint;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.ResourceAggregate;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
@@ -33,6 +30,7 @@ import org.apache.aurora.gen.TaskConstraint;
 import org.apache.aurora.gen.ValueConstraint;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.configuration.ConfigurationManager;
 import org.apache.aurora.scheduler.storage.Storage;
 import org.apache.aurora.scheduler.storage.entities.IJobKey;
@@ -167,28 +165,21 @@ public class ResourceCounterTest {
       ScheduleStatus status,
       Optional<String> dedicated) {
 
-    TaskConfig task = new TaskConfig()
-        .setJob(new JobKey(role, "test", job))
-        .setOwner(new Identity().setRole(role))
-        .setEnvironment("test")
-        .setJobName(job)
+    ScheduledTask task = TaskTestUtil.makeTask(id, JobKeys.from(role, "test", job)).newBuilder();
+    TaskConfig config = task.getAssignedTask().getTask()
         .setNumCpus(numCpus)
         .setRamMb(ramMb)
         .setDiskMb(diskMb)
         .setProduction(production);
 
     if (dedicated.isPresent()) {
-      task.addToConstraints(new Constraint(
+      config.addToConstraints(new Constraint(
           ConfigurationManager.DEDICATED_ATTRIBUTE,
           TaskConstraint.value(new ValueConstraint(false, ImmutableSet.of(dedicated.get())))));
     }
 
-    return IScheduledTask.build(new ScheduledTask()
-        .setStatus(status)
-        .setAssignedTask(
-            new AssignedTask()
-                .setTaskId(id)
-                .setTask(task)));
+    task.setStatus(status);
+    return IScheduledTask.build(task);
   }
 
   private void insertTasks(final IScheduledTask... tasks) {

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/storage/StorageBackfillTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/StorageBackfillTest.java b/src/test/java/org/apache/aurora/scheduler/storage/StorageBackfillTest.java
index 90a05ba..254b231 100644
--- a/src/test/java/org/apache/aurora/scheduler/storage/StorageBackfillTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/storage/StorageBackfillTest.java
@@ -15,50 +15,34 @@ package org.apache.aurora.scheduler.storage;
 
 import java.util.Set;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Maps;
-import com.google.common.collect.Sets;
 
-import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.Constraint;
 import org.apache.aurora.gen.Container;
-import org.apache.aurora.gen.ExecutorConfig;
-import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.JobConfiguration;
-import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.MesosContainer;
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.TaskConfig;
-import org.apache.aurora.gen.TaskEvent;
-import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.configuration.ConfigurationManager;
 import org.apache.aurora.scheduler.configuration.SanitizedConfiguration;
 import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
-import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
-import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
 import org.apache.aurora.scheduler.storage.mem.MemStorage;
 import org.junit.Before;
 import org.junit.Test;
 
-import static org.apache.aurora.gen.ScheduleStatus.PENDING;
-import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
 import static org.junit.Assert.assertEquals;
 
+/**
+ * TODO(wfarner): Data store constraints (not null, valid relations) have made this test of little
+ * value other than for coverage.  Rethink.
+ */
 public class StorageBackfillTest {
-  private static final String ROLE = "Test_Role_A";
-  private static final String USER = "Test_User_A";
-  private static final Identity OWNER = new Identity(ROLE, USER);
-  private static final String ENV = "Test_Env";
-  private static final String JOB_NAME = "Test_Job";
-  private static final IJobKey JOB_KEY = JobKeys.from(ROLE, ENV, JOB_NAME);
-  private static final int ONE_GB = 1024;
-  private static final String TASK_ID = "task_id";
-  private static final ExecutorConfig EXECUTOR_CONFIG =
-      new ExecutorConfig("AuroraExecutor", "executorConfig");
+  // By default, only the mesos container is allowed by ConfigurationManager.
+  private static final IScheduledTask TASK =
+      setMesosContainer(TaskTestUtil.makeTask("task_id", TaskTestUtil.JOB));
 
   private Storage storage;
 
@@ -68,40 +52,14 @@ public class StorageBackfillTest {
   }
 
   @Test
-  public void testLoadTasksFromStorage() throws Exception {
-    final TaskConfig storedTask = defaultTask();
-
-    storage.write(new Storage.MutateWork.NoResult.Quiet() {
-      @Override
-      protected void execute(Storage.MutableStoreProvider storeProvider) {
-        storeProvider.getUnsafeTaskStore().saveTasks(ImmutableSet.of(
-            IScheduledTask.build(new ScheduledTask()
-                .setStatus(PENDING)
-                .setTaskEvents(ImmutableList.of(new TaskEvent(100, PENDING)))
-                .setAssignedTask(new AssignedTask()
-                    .setTaskId(TASK_ID)
-                    .setInstanceId(0)
-                    .setTask(storedTask)))));
-      }
-    });
-
-    backfill();
-
-    // Since task fields are backfilled with defaults, additional flags should be filled.
-    ITaskConfig expected = ITaskConfig.build(new TaskConfig(storedTask)
-        .setJob(JOB_KEY.newBuilder())
-        .setProduction(false)
-        .setMaxTaskFailures(1)
-        .setExecutorConfig(EXECUTOR_CONFIG)
-        .setContainer(Container.mesos(new MesosContainer()))
-        .setConstraints(ImmutableSet.of(ConfigurationManager.hostLimitConstraint(1))));
-
-    assertEquals(expected, getTask(TASK_ID).getAssignedTask().getTask());
-  }
-
-  @Test
   public void testJobConfigurationBackfill() throws Exception {
-    final JobConfiguration config = makeJobConfig(JOB_KEY, defaultTask(), 1);
+    TaskConfig task = TASK.getAssignedTask().getTask().newBuilder();
+    final JobConfiguration config = new JobConfiguration()
+        .setOwner(task.getOwner())
+        .setKey(task.getJob())
+        .setInstanceCount(1)
+        .setTaskConfig(task);
+
     SanitizedConfiguration expected =
         SanitizedConfiguration.fromUnsanitized(IJobConfiguration.build(config));
 
@@ -128,25 +86,11 @@ public class StorageBackfillTest {
   }
 
   @Test
-  public void testBackfillTaskJob() throws Exception {
-    TaskConfig task = defaultTask();
-    ConfigurationManager.applyDefaultsIfUnset(task);
-    task.unsetJob();
-
-    IScheduledTask noJobKey = IScheduledTask.build(new ScheduledTask()
-        .setStatus(RUNNING)
-        .setAssignedTask(new AssignedTask()
-            .setInstanceId(0)
-            .setTaskId("nojobkey")
-            .setTask(task)));
-    IScheduledTask nullJobKeyFields = IScheduledTask.build(new ScheduledTask()
-        .setStatus(RUNNING)
-        .setAssignedTask(new AssignedTask()
-            .setInstanceId(1)
-            .setTaskId("nulled_fields")
-            .setTask(task.setJob(new JobKey()))));
-
-    final Set<IScheduledTask> backfilledTasks = ImmutableSet.of(noJobKey, nullJobKeyFields);
+  public void testBackfillTask() throws Exception {
+    ScheduledTask task = TASK.newBuilder();
+    ConfigurationManager.applyDefaultsIfUnset(task.getAssignedTask().getTask());
+
+    final Set<IScheduledTask> backfilledTasks = ImmutableSet.of(IScheduledTask.build(task));
     storage.write(new Storage.MutateWork.NoResult.Quiet() {
       @Override
       protected void execute(Storage.MutableStoreProvider storeProvider) {
@@ -155,18 +99,9 @@ public class StorageBackfillTest {
     });
 
     backfill();
-    ScheduledTask noJobKeyBackfilled = noJobKey.newBuilder();
-    noJobKeyBackfilled.getAssignedTask().getTask()
-        .setJob(new JobKey(OWNER.getRole(), ENV, JOB_NAME));
-
-    ScheduledTask nullJobKeyFieldsBackfilled = nullJobKeyFields.newBuilder();
-    nullJobKeyFieldsBackfilled.getAssignedTask().getTask()
-        .setJob(new JobKey(OWNER.getRole(), ENV, JOB_NAME));
 
     assertEquals(
-        ImmutableSet.of(
-            IScheduledTask.build(noJobKeyBackfilled),
-            IScheduledTask.build(nullJobKeyFieldsBackfilled)),
+        ImmutableSet.of(IScheduledTask.build(task)),
         Storage.Util.fetchTasks(storage, Query.unscoped()));
   }
 
@@ -179,34 +114,9 @@ public class StorageBackfillTest {
     });
   }
 
-  private static JobConfiguration makeJobConfig(IJobKey jobKey, TaskConfig task, int numTasks) {
-    return new JobConfiguration()
-        .setOwner(OWNER)
-        .setKey(jobKey.newBuilder())
-        .setInstanceCount(numTasks)
-        .setTaskConfig(new TaskConfig(task)
-            .setOwner(OWNER)
-            .setEnvironment(jobKey.getEnvironment())
-            .setJobName(jobKey.getName()));
-  }
-
-  private static TaskConfig defaultTask() {
-    return new TaskConfig()
-        .setOwner(OWNER)
-        .setJobName(JOB_NAME)
-        .setEnvironment(ENV)
-        .setNumCpus(1.0)
-        .setRamMb(ONE_GB)
-        .setDiskMb(500)
-        .setExecutorConfig(EXECUTOR_CONFIG)
-        .setRequestedPorts(Sets.<String>newHashSet())
-        .setConstraints(Sets.<Constraint>newHashSet())
-        .setTaskLinks(Maps.<String, String>newHashMap());
-  }
-
-  private IScheduledTask getTask(String taskId) {
-    return Iterables.getOnlyElement(Storage.Util.fetchTasks(
-        storage,
-        Query.taskScoped(taskId)));
+  private static IScheduledTask setMesosContainer(IScheduledTask task) {
+    ScheduledTask builder = task.newBuilder();
+    builder.getAssignedTask().getTask().setContainer(Container.mesos(new MesosContainer()));
+    return IScheduledTask.build(builder);
   }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/449a835e/src/test/java/org/apache/aurora/scheduler/storage/backup/RecoveryTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/backup/RecoveryTest.java b/src/test/java/org/apache/aurora/scheduler/storage/backup/RecoveryTest.java
index 5853c03..09e4208 100644
--- a/src/test/java/org/apache/aurora/scheduler/storage/backup/RecoveryTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/storage/backup/RecoveryTest.java
@@ -25,18 +25,15 @@ import com.twitter.common.quantity.Time;
 import com.twitter.common.testing.easymock.EasyMockTest;
 import com.twitter.common.util.testing.FakeClock;
 
-import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.HostAttributes;
-import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.Lock;
-import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.gen.storage.QuotaConfiguration;
 import org.apache.aurora.gen.storage.SchedulerMetadata;
 import org.apache.aurora.gen.storage.Snapshot;
 import org.apache.aurora.gen.storage.StoredCronJob;
 import org.apache.aurora.gen.storage.StoredJobUpdateDetails;
 import org.apache.aurora.scheduler.base.Query;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.storage.DistributedSnapshotStore;
 import org.apache.aurora.scheduler.storage.SnapshotStore;
@@ -62,8 +59,8 @@ import static org.junit.Assert.assertEquals;
 public class RecoveryTest extends EasyMockTest {
 
   private static final Amount<Long, Time> INTERVAL = Amount.of(1L, Time.HOURS);
-  private static final ScheduledTask TASK1 = makeTask("task1");
-  private static final ScheduledTask TASK2 = makeTask("task2");
+  private static final IScheduledTask TASK1 = TaskTestUtil.makeTask("task1", TaskTestUtil.JOB);
+  private static final IScheduledTask TASK2 = TaskTestUtil.makeTask("task2", TaskTestUtil.JOB);
   private static final Snapshot SNAPSHOT1 = makeSnapshot(TASK1, TASK2);
 
   private SnapshotStore<Snapshot> snapshotStore;
@@ -129,7 +126,7 @@ public class RecoveryTest extends EasyMockTest {
   @Test
   public void testModifySnapshotBeforeCommit() throws Exception {
     expect(snapshotStore.createSnapshot()).andReturn(SNAPSHOT1);
-    Snapshot modified = SNAPSHOT1.deepCopy().setTasks(ImmutableSet.of(TASK1));
+    Snapshot modified = SNAPSHOT1.deepCopy().setTasks(ImmutableSet.of(TASK1.newBuilder()));
     Capture<MutateWork<Object, Exception>> transaction = createCapture();
     expect(primaryStorage.write(capture(transaction))).andReturn(null);
     distributedStore.persist(modified);
@@ -164,24 +161,14 @@ public class RecoveryTest extends EasyMockTest {
     recovery.commit();
   }
 
-  private static Snapshot makeSnapshot(ScheduledTask... tasks) {
+  private static Snapshot makeSnapshot(IScheduledTask... tasks) {
     return new Snapshot()
         .setHostAttributes(ImmutableSet.<HostAttributes>of())
         .setCronJobs(ImmutableSet.<StoredCronJob>of())
         .setSchedulerMetadata(new SchedulerMetadata().setVersion(CURRENT_API_VERSION))
         .setQuotaConfigurations(ImmutableSet.<QuotaConfiguration>of())
-        .setTasks(ImmutableSet.<ScheduledTask>builder().add(tasks).build())
+        .setTasks(IScheduledTask.toBuildersSet(ImmutableSet.copyOf(tasks)))
         .setLocks(ImmutableSet.<Lock>of())
         .setJobUpdateDetails(ImmutableSet.<StoredJobUpdateDetails>of());
   }
-
-  private static ScheduledTask makeTask(String taskId) {
-    return new ScheduledTask().setAssignedTask(
-        new AssignedTask()
-            .setTaskId(taskId)
-            .setTask(new TaskConfig()
-                .setJobName("job-" + taskId)
-                .setEnvironment("test")
-                .setOwner(new Identity().setRole("role-" + taskId).setUser("user-" + taskId))));
-  }
 }


Mime
View raw message