You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@aurora.apache.org by wf...@apache.org on 2016/01/23 02:20:46 UTC

aurora git commit: Remove most direct uses of deprecated TaskConfig fields.

Repository: aurora
Updated Branches:
  refs/heads/master 02beb1f12 -> 5684c64c4


Remove most direct uses of deprecated TaskConfig fields.

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


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

Branch: refs/heads/master
Commit: 5684c64c46fcf74d3da3b81f7a3e0a10b5eddaca
Parents: 02beb1f
Author: Bill Farner <wf...@apache.org>
Authored: Fri Jan 22 17:14:28 2016 -0800
Committer: Bill Farner <wf...@apache.org>
Committed: Fri Jan 22 17:14:28 2016 -0800

----------------------------------------------------------------------
 .../aurora/scheduler/http/Utilization.java      |   2 +-
 .../scheduler/state/MaintenanceController.java  |  39 +++---
 .../aurora/scheduler/storage/TaskStore.java     |   4 +-
 .../thrift/SchedulerThriftInterface.java        |  13 +-
 .../aurora/codec/ThriftBinaryCodecTest.java     |  21 ++--
 .../apache/aurora/scheduler/TaskVarsTest.java   |  42 +++----
 .../aurora/scheduler/app/SchedulerIT.java       |  37 +++---
 .../scheduler/cron/quartz/QuartzTestUtil.java   |  23 +---
 .../filter/SchedulingFilterImplTest.java        | 120 ++++++++-----------
 .../aurora/scheduler/http/api/ApiBetaTest.java  |  24 +---
 .../mesos/MesosTaskFactoryImplTest.java         |  17 +--
 .../pruning/TaskHistoryPrunerTest.java          |  35 ++----
 .../scheduler/quota/QuotaManagerImplTest.java   |  29 ++---
 .../RescheduleCalculatorImplTest.java           |  15 +--
 .../scheduler/sla/MetricCalculatorTest.java     |   1 -
 .../aurora/scheduler/sla/SlaTestUtil.java       |  25 ++--
 .../scheduler/state/LockManagerImplTest.java    |   4 +-
 .../state/MaintenanceControllerImplTest.java    |  48 ++++----
 .../scheduler/state/TaskStateMachineTest.java   |  20 ++--
 .../thrift/SchedulerThriftInterfaceTest.java    |  26 ++--
 .../updater/JobUpdateEventSubscriberTest.java   |  23 +---
 21 files changed, 216 insertions(+), 352 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/main/java/org/apache/aurora/scheduler/http/Utilization.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/http/Utilization.java b/src/main/java/org/apache/aurora/scheduler/http/Utilization.java
index 440946b..4a609e3 100644
--- a/src/main/java/org/apache/aurora/scheduler/http/Utilization.java
+++ b/src/main/java/org/apache/aurora/scheduler/http/Utilization.java
@@ -206,7 +206,7 @@ public class Utilization {
       @PathParam("role") String role) {
 
     MetricType type = getTypeByName(metric);
-    Function<ITaskConfig, Display> toKey = task -> new Display(task.getJobName(), null);
+    Function<ITaskConfig, Display> toKey = task -> new Display(task.getJob().getName(), null);
     Map<Display, Metric> byJob =
         counter.computeAggregates(Query.roleScoped(role).active(), type.filter, toKey);
     return Response.ok(fillTemplate(byJob)).build();

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java b/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
index 3770056..3c7cda0 100644
--- a/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
+++ b/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
@@ -39,6 +39,7 @@ import org.apache.aurora.scheduler.storage.Storage;
 import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
 import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
+import org.apache.aurora.scheduler.storage.entities.IHostStatus;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -66,7 +67,7 @@ public interface MaintenanceController {
    * @param hosts Hosts to put into maintenance mode.
    * @return The adjusted state of the hosts.
    */
-  Set<HostStatus> startMaintenance(Set<String> hosts);
+  Set<IHostStatus> startMaintenance(Set<String> hosts);
 
   /**
    * Initiate a drain of all active tasks on {@code hosts}.
@@ -75,7 +76,7 @@ public interface MaintenanceController {
    * @return The adjusted state of the hosts.  Hosts without any active tasks will be immediately
    *         moved to DRAINED.
    */
-  Set<HostStatus> drain(Set<String> hosts);
+  Set<IHostStatus> drain(Set<String> hosts);
 
   /**
    * Fetches the current maintenance mode of {$code host}.
@@ -91,7 +92,7 @@ public interface MaintenanceController {
    * @param hosts Hosts to fetch state for.
    * @return The state of the hosts.
    */
-  Set<HostStatus> getStatus(Set<String> hosts);
+  Set<IHostStatus> getStatus(Set<String> hosts);
 
   /**
    * Moves {@code hosts} out of maintenance mode, returning them to mode NONE.
@@ -99,7 +100,7 @@ public interface MaintenanceController {
    * @param hosts Hosts to move out of maintenance mode.
    * @return The adjusted state of the hosts.
    */
-  Set<HostStatus> endMaintenance(Set<String> hosts);
+  Set<IHostStatus> endMaintenance(Set<String> hosts);
 
   class MaintenanceControllerImpl implements MaintenanceController, EventSubscriber {
     private static final Logger LOG = LoggerFactory.getLogger(MaintenanceControllerImpl.class);
@@ -112,7 +113,7 @@ public interface MaintenanceController {
       this.stateManager = requireNonNull(stateManager);
     }
 
-    private Set<HostStatus> watchDrainingTasks(MutableStoreProvider store, Set<String> hosts) {
+    private Set<IHostStatus> watchDrainingTasks(MutableStoreProvider store, Set<String> hosts) {
       LOG.info("Hosts to drain: " + hosts);
       Set<String> emptyHosts = Sets.newHashSet();
       for (String host : hosts) {
@@ -137,7 +138,7 @@ public interface MaintenanceController {
         }
       }
 
-      return ImmutableSet.<HostStatus>builder()
+      return ImmutableSet.<IHostStatus>builder()
           .addAll(setMaintenanceMode(store, emptyHosts, DRAINED))
           .addAll(setMaintenanceMode(store, Sets.difference(hosts, emptyHosts), DRAINING))
           .build();
@@ -172,7 +173,7 @@ public interface MaintenanceController {
     }
 
     @Override
-    public Set<HostStatus> startMaintenance(final Set<String> hosts) {
+    public Set<IHostStatus> startMaintenance(Set<String> hosts) {
       return storage.write(
           storeProvider -> setMaintenanceMode(storeProvider, hosts, MaintenanceMode.SCHEDULED));
     }
@@ -182,18 +183,18 @@ public interface MaintenanceController {
         Optional.of("Draining machine for maintenance.");
 
     @Override
-    public Set<HostStatus> drain(final Set<String> hosts) {
+    public Set<IHostStatus> drain(Set<String> hosts) {
       return storage.write(store -> watchDrainingTasks(store, hosts));
     }
 
     private static final Function<IHostAttributes, String> HOST_NAME =
         IHostAttributes::getHost;
 
-    private static final Function<IHostAttributes, HostStatus> ATTRS_TO_STATUS =
-        attributes -> new HostStatus().setHost(attributes.getHost()).setMode(attributes.getMode());
+    private static final Function<IHostAttributes, IHostStatus> ATTRS_TO_STATUS =
+        attributes -> IHostStatus.build(
+            new HostStatus().setHost(attributes.getHost()).setMode(attributes.getMode()));
 
-    private static final Function<HostStatus, MaintenanceMode> GET_MODE =
-        HostStatus::getMode;
+    private static final Function<IHostStatus, MaintenanceMode> GET_MODE = IHostStatus::getMode;
 
     @Override
     public MaintenanceMode getMode(final String host) {
@@ -204,37 +205,37 @@ public interface MaintenanceController {
     }
 
     @Override
-    public Set<HostStatus> getStatus(final Set<String> hosts) {
+    public Set<IHostStatus> getStatus(final Set<String> hosts) {
       return storage.read(storeProvider -> {
         // Warning - this is filtering _all_ host attributes.  If using this to frequently query
         // for a small set of hosts, a getHostAttributes variant should be added.
         return FluentIterable.from(storeProvider.getAttributeStore().getHostAttributes())
             .filter(Predicates.compose(Predicates.in(hosts), HOST_NAME))
-            .transform(ATTRS_TO_STATUS).toSet();
+            .transform(ATTRS_TO_STATUS)
+            .toSet();
       });
     }
 
     @Override
-    public Set<HostStatus> endMaintenance(final Set<String> hosts) {
+    public Set<IHostStatus> endMaintenance(final Set<String> hosts) {
       return storage.write(
           storeProvider -> setMaintenanceMode(storeProvider, hosts, MaintenanceMode.NONE));
     }
 
-    private Set<HostStatus> setMaintenanceMode(
+    private Set<IHostStatus> setMaintenanceMode(
         MutableStoreProvider storeProvider,
         Set<String> hosts,
         MaintenanceMode mode) {
 
       AttributeStore.Mutable store = storeProvider.getAttributeStore();
-      ImmutableSet.Builder<HostStatus> statuses = ImmutableSet.builder();
+      ImmutableSet.Builder<IHostStatus> statuses = ImmutableSet.builder();
       for (String host : hosts) {
         LOG.info("Setting maintenance mode to {} for host {}", mode, host);
         Optional<IHostAttributes> toSave = AttributeStore.Util.mergeMode(store, host, mode);
         if (toSave.isPresent()) {
           store.saveHostAttributes(toSave.get());
           LOG.info("Updated host attributes: " + toSave.get());
-          HostStatus status = new HostStatus().setHost(host).setMode(mode);
-          statuses.add(status);
+          statuses.add(IHostStatus.build(new HostStatus().setHost(host).setMode(mode)));
         }
       }
       return statuses.build();

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java b/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
index 3ab2d54..ac0bb37 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
@@ -133,10 +133,10 @@ public interface TaskStore {
           return false;
         }
         if (query.getEnvironment() != null
-            && !query.getEnvironment().equals(config.getEnvironment())) {
+            && !query.getEnvironment().equals(config.getJob().getEnvironment())) {
           return false;
         }
-        if (query.getJobName() != null && !query.getJobName().equals(config.getJobName())) {
+        if (query.getJobName() != null && !query.getJobName().equals(config.getJob().getName())) {
           return false;
         }
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
index d1e3c83..662cae1 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
@@ -95,6 +95,7 @@ import org.apache.aurora.scheduler.storage.backup.Recovery;
 import org.apache.aurora.scheduler.storage.backup.StorageBackup;
 import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
 import org.apache.aurora.scheduler.storage.entities.IConfigRewrite;
+import org.apache.aurora.scheduler.storage.entities.IHostStatus;
 import org.apache.aurora.scheduler.storage.entities.IInstanceConfigRewrite;
 import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
 import org.apache.aurora.scheduler.storage.entities.IJobConfigRewrite;
@@ -543,26 +544,30 @@ class SchedulerThriftInterface implements AnnotatedAuroraAdmin {
   public Response startMaintenance(Hosts hosts) {
     return ok(Result.startMaintenanceResult(
         new StartMaintenanceResult()
-            .setStatuses(maintenance.startMaintenance(hosts.getHostNames()))));
+            .setStatuses(IHostStatus.toBuildersSet(
+                maintenance.startMaintenance(hosts.getHostNames())))));
   }
 
   @Override
   public Response drainHosts(Hosts hosts) {
     return ok(Result.drainHostsResult(
-        new DrainHostsResult().setStatuses(maintenance.drain(hosts.getHostNames()))));
+        new DrainHostsResult().setStatuses(IHostStatus.toBuildersSet(
+            maintenance.drain(hosts.getHostNames())))));
   }
 
   @Override
   public Response maintenanceStatus(Hosts hosts) {
     return ok(Result.maintenanceStatusResult(
-        new MaintenanceStatusResult().setStatuses(maintenance.getStatus(hosts.getHostNames()))));
+        new MaintenanceStatusResult().setStatuses(IHostStatus.toBuildersSet(
+            maintenance.getStatus(hosts.getHostNames())))));
   }
 
   @Override
   public Response endMaintenance(Hosts hosts) {
     return ok(Result.endMaintenanceResult(
         new EndMaintenanceResult()
-            .setStatuses(maintenance.endMaintenance(hosts.getHostNames()))));
+            .setStatuses(IHostStatus.toBuildersSet(
+                maintenance.endMaintenance(hosts.getHostNames())))));
   }
 
   @Override

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/codec/ThriftBinaryCodecTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/codec/ThriftBinaryCodecTest.java b/src/test/java/org/apache/aurora/codec/ThriftBinaryCodecTest.java
index ebb4f9a..999ef06 100644
--- a/src/test/java/org/apache/aurora/codec/ThriftBinaryCodecTest.java
+++ b/src/test/java/org/apache/aurora/codec/ThriftBinaryCodecTest.java
@@ -14,7 +14,8 @@
 package org.apache.aurora.codec;
 
 import org.apache.aurora.codec.ThriftBinaryCodec.CodingException;
-import org.apache.aurora.gen.Identity;
+import org.apache.aurora.gen.ScheduledTask;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
@@ -24,21 +25,23 @@ public class ThriftBinaryCodecTest {
 
   @Test
   public void testRoundTrip() throws CodingException {
-    Identity original = new Identity("mesos", "jack");
+    ScheduledTask original = TaskTestUtil.makeTask("id", TaskTestUtil.JOB).newBuilder();
     assertEquals(original,
-        ThriftBinaryCodec.decode(Identity.class, ThriftBinaryCodec.encode(original)));
+        ThriftBinaryCodec.decode(ScheduledTask.class, ThriftBinaryCodec.encode(original)));
   }
 
   @Test
   public void testRoundTripNull() throws CodingException {
-    assertNull(ThriftBinaryCodec.decode(Identity.class, ThriftBinaryCodec.encode(null)));
+    assertNull(ThriftBinaryCodec.decode(ScheduledTask.class, ThriftBinaryCodec.encode(null)));
   }
 
   @Test
   public void testRoundTripNonNull() throws CodingException {
-    Identity original = new Identity("mesos", "jill");
+    ScheduledTask original = TaskTestUtil.makeTask("id", TaskTestUtil.JOB).newBuilder();
     assertEquals(original,
-        ThriftBinaryCodec.decodeNonNull(Identity.class, ThriftBinaryCodec.encodeNonNull(original)));
+        ThriftBinaryCodec.decodeNonNull(
+            ScheduledTask.class,
+            ThriftBinaryCodec.encodeNonNull(original)));
   }
 
   @Test(expected = NullPointerException.class)
@@ -48,16 +51,16 @@ public class ThriftBinaryCodecTest {
 
   @Test(expected = NullPointerException.class)
   public void testDecodeNonNull() throws CodingException {
-    ThriftBinaryCodec.decodeNonNull(Identity.class, null);
+    ThriftBinaryCodec.decodeNonNull(ScheduledTask.class, null);
   }
 
   @Test
   public void testInflateDeflateRoundTrip() throws CodingException {
-    Identity original = new Identity("aurora", "jsmith");
+    ScheduledTask original = TaskTestUtil.makeTask("id", TaskTestUtil.JOB).newBuilder();
 
     byte[] deflated = ThriftBinaryCodec.deflateNonNull(original);
 
-    Identity inflated = ThriftBinaryCodec.inflateNonNull(Identity.class, deflated);
+    ScheduledTask inflated = ThriftBinaryCodec.inflateNonNull(ScheduledTask.class, deflated);
 
     assertEquals(original, inflated);
   }

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/TaskVarsTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/TaskVarsTest.java b/src/test/java/org/apache/aurora/scheduler/TaskVarsTest.java
index 0a9dfe3..05cd78f 100644
--- a/src/test/java/org/apache/aurora/scheduler/TaskVarsTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/TaskVarsTest.java
@@ -23,15 +23,13 @@ import com.google.common.collect.Maps;
 
 import org.apache.aurora.common.stats.StatsProvider;
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
-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.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
+import org.apache.aurora.scheduler.base.JobKeys;
 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;
@@ -39,6 +37,7 @@ import org.apache.aurora.scheduler.events.PubsubEvent.TasksDeleted;
 import org.apache.aurora.scheduler.filter.SchedulingFilter.Veto;
 import org.apache.aurora.scheduler.filter.SchedulingFilter.VetoGroup;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
+import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
 import org.easymock.EasyMock;
@@ -63,11 +62,8 @@ import static org.junit.Assert.assertNotNull;
 
 public class TaskVarsTest extends EasyMockTest {
 
-  private static final String ROLE_A = "role_a";
-  private static final String JOB_A = "job_a";
-  private static final String JOB_B = "job_b";
-  private static final String TASK_ID = "task_id";
-  private static final String ENV = "test";
+  private static final IJobKey JOB_A = JobKeys.from("role_a", "test", "job_a");
+  private static final IJobKey JOB_B = JobKeys.from("role_a", "test", "job_b");
 
   private static final String STATIC_COUNTER = VETO_GROUPS_TO_COUNTERS.get(VetoGroup.STATIC);
   private static final String DYNAMIC_COUNTER = VETO_GROUPS_TO_COUNTERS.get(VetoGroup.DYNAMIC);
@@ -102,7 +98,9 @@ public class TaskVarsTest extends EasyMockTest {
     expect(provider.makeGauge(EasyMock.eq(name), EasyMock.<Supplier<Long>>anyObject()))
         .andAnswer(() -> {
           assertFalse(globalCounters.containsKey(name));
-          globalCounters.put(name, (Supplier<Long>) EasyMock.getCurrentArguments()[1]);
+          @SuppressWarnings("unchecked")
+          Supplier<Long> varSupplier = (Supplier<Long>) EasyMock.getCurrentArguments()[1];
+          globalCounters.put(name, varSupplier);
           return null;
         });
   }
@@ -132,16 +130,9 @@ public class TaskVarsTest extends EasyMockTest {
     vars.startAsync().awaitRunning();
   }
 
-  private IScheduledTask makeTask(String job, ScheduleStatus status, String host) {
-    ScheduledTask task = new ScheduledTask()
-        .setStatus(status)
-        .setAssignedTask(new AssignedTask()
-            .setTaskId(TASK_ID)
-            .setTask(new TaskConfig()
-                .setJob(new JobKey(ROLE_A, ENV, job))
-                .setJobName(job)
-                .setEnvironment(ENV)
-                .setOwner(new Identity(ROLE_A, ROLE_A + "-user"))));
+  private IScheduledTask makeTask(IJobKey job, ScheduleStatus status, String host) {
+    ScheduledTask task = TaskTestUtil.makeTask("task_id", job).newBuilder()
+        .setStatus(status);
     if (Tasks.SLAVE_ASSIGNED_STATES.contains(status) || Tasks.isTerminated(status)) {
       task.getAssignedTask().setSlaveHost(host);
     }
@@ -149,7 +140,7 @@ public class TaskVarsTest extends EasyMockTest {
     return IScheduledTask.build(task);
   }
 
-  private IScheduledTask makeTask(String job, ScheduleStatus status) {
+  private IScheduledTask makeTask(IJobKey job, ScheduleStatus status) {
     return makeTask(job, status, "hostA");
   }
 
@@ -306,10 +297,11 @@ public class TaskVarsTest extends EasyMockTest {
     expectStatExport(rackStatName("rackA"));
     expectStatExport(rackStatName("rackB"));
 
-    IScheduledTask a = makeTask("jobA", RUNNING, "host1");
-    IScheduledTask b = makeTask("jobB", RUNNING, "host2");
-    IScheduledTask c = makeTask("jobD", RUNNING, "host3");
-    IScheduledTask d = makeTask("jobD", RUNNING, "host1");
+    IScheduledTask a = makeTask(JOB_A, RUNNING, "host1");
+    IScheduledTask b = makeTask(JOB_B, RUNNING, "host2");
+    IJobKey jobD = JobKeys.from(JOB_A.getRole(), JOB_A.getEnvironment(), "jobD");
+    IScheduledTask c = makeTask(jobD, RUNNING, "host3");
+    IScheduledTask d = makeTask(jobD, RUNNING, "host1");
 
     expectStatExport(jobStatName(a, LOST), untrackedProvider);
     expectStatExport(jobStatName(b, LOST), untrackedProvider);

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/app/SchedulerIT.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/app/SchedulerIT.java b/src/test/java/org/apache/aurora/scheduler/app/SchedulerIT.java
index 83f4b7a..39916f5 100644
--- a/src/test/java/org/apache/aurora/scheduler/app/SchedulerIT.java
+++ b/src/test/java/org/apache/aurora/scheduler/app/SchedulerIT.java
@@ -49,14 +49,9 @@ import org.apache.aurora.common.zookeeper.ServerSetImpl;
 import org.apache.aurora.common.zookeeper.ZooKeeperClient;
 import org.apache.aurora.common.zookeeper.ZooKeeperClient.Credentials;
 import org.apache.aurora.common.zookeeper.testing.BaseZooKeeperTest;
-import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.ServerInfo;
-import org.apache.aurora.gen.TaskConfig;
-import org.apache.aurora.gen.TaskEvent;
 import org.apache.aurora.gen.storage.LogEntry;
 import org.apache.aurora.gen.storage.Op;
 import org.apache.aurora.gen.storage.SaveFrameworkId;
@@ -66,6 +61,7 @@ import org.apache.aurora.gen.storage.Transaction;
 import org.apache.aurora.gen.storage.storageConstants;
 import org.apache.aurora.scheduler.AppStartup;
 import org.apache.aurora.scheduler.ResourceSlot;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.configuration.executor.ExecutorSettings;
 import org.apache.aurora.scheduler.log.Log;
 import org.apache.aurora.scheduler.log.Log.Entry;
@@ -75,6 +71,7 @@ import org.apache.aurora.scheduler.mesos.DriverFactory;
 import org.apache.aurora.scheduler.mesos.DriverSettings;
 import org.apache.aurora.scheduler.mesos.TestExecutorSettings;
 import org.apache.aurora.scheduler.storage.backup.BackupModule;
+import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.entities.IServerInfo;
 import org.apache.aurora.scheduler.storage.log.EntrySerializer;
 import org.apache.aurora.scheduler.storage.log.LogStorageModule;
@@ -272,19 +269,14 @@ public class SchedulerIT extends BaseZooKeeperTest {
         });
   }
 
-  private static ScheduledTask makeTask(String id, ScheduleStatus status) {
-    return new ScheduledTask()
-        .setStatus(status)
-        .setTaskEvents(ImmutableList.of(new TaskEvent(100, status)))
-        .setAssignedTask(new AssignedTask()
-            .setSlaveId("slaveId")
-            .setTaskId(id)
-            .setTask(new TaskConfig()
-                .setJob(new JobKey("role-" + id, "test", "job-" + id))
-                .setJobName("job-" + id)
-                .setEnvironment("test")
-                .setExecutorConfig(new org.apache.aurora.gen.ExecutorConfig("AuroraExecutor", ""))
-                .setOwner(new Identity("role-" + id, "user-" + id))));
+  private static IScheduledTask makeTask(String id, ScheduleStatus status) {
+    ScheduledTask builder = TaskTestUtil.addStateTransition(
+        TaskTestUtil.makeTask(id, TaskTestUtil.JOB),
+        status,
+        100)
+        .newBuilder();
+    builder.getAssignedTask().setSlaveId("slave-id");
+    return IScheduledTask.build(builder);
   }
 
   @Test
@@ -297,12 +289,13 @@ public class SchedulerIT extends BaseZooKeeperTest {
         eq(SETTINGS.getMasterUri())))
         .andReturn(driver).anyTimes();
 
-    ScheduledTask snapshotTask = makeTask("snapshotTask", ScheduleStatus.ASSIGNED);
-    ScheduledTask transactionTask = makeTask("transactionTask", ScheduleStatus.RUNNING);
+    IScheduledTask snapshotTask = makeTask("snapshotTask", ScheduleStatus.ASSIGNED);
+    IScheduledTask transactionTask = makeTask("transactionTask", ScheduleStatus.RUNNING);
     Iterable<Entry> recoveredEntries = toEntries(
-        LogEntry.snapshot(new Snapshot().setTasks(ImmutableSet.of(snapshotTask))),
+        LogEntry.snapshot(new Snapshot().setTasks(ImmutableSet.of(snapshotTask.newBuilder()))),
         LogEntry.transaction(new Transaction(
-            ImmutableList.of(Op.saveTasks(new SaveTasks(ImmutableSet.of(transactionTask)))),
+            ImmutableList.of(Op.saveTasks(
+                new SaveTasks(ImmutableSet.of(transactionTask.newBuilder())))),
             storageConstants.CURRENT_SCHEMA_VERSION)));
 
     expect(log.open()).andReturn(logStream);

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/cron/quartz/QuartzTestUtil.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/cron/quartz/QuartzTestUtil.java b/src/test/java/org/apache/aurora/scheduler/cron/quartz/QuartzTestUtil.java
index f062497..ef9aae1 100644
--- a/src/test/java/org/apache/aurora/scheduler/cron/quartz/QuartzTestUtil.java
+++ b/src/test/java/org/apache/aurora/scheduler/cron/quartz/QuartzTestUtil.java
@@ -14,14 +14,12 @@
 package org.apache.aurora.scheduler.cron.quartz;
 
 import com.google.common.base.Throwables;
-import com.google.common.collect.ImmutableSet;
 
+import org.apache.aurora.gen.Container;
 import org.apache.aurora.gen.CronCollisionPolicy;
-import org.apache.aurora.gen.ExecutorConfig;
 import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.JobConfiguration;
-import org.apache.aurora.gen.Metadata;
-import org.apache.aurora.gen.TaskConfig;
+import org.apache.aurora.gen.MesosContainer;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.configuration.ConfigurationManager;
@@ -42,19 +40,10 @@ final class QuartzTestUtil {
           .setInstanceCount(10)
           .setOwner(new Identity("role", "user"))
           .setKey(AURORA_JOB_KEY.newBuilder())
-          .setTaskConfig(new TaskConfig()
-              .setJob(AURORA_JOB_KEY.newBuilder())
-              .setOwner(new Identity("role", "user"))
-              .setJobName(AURORA_JOB_KEY.getName())
-              .setEnvironment(AURORA_JOB_KEY.getEnvironment())
-              .setDiskMb(3)
-              .setRamMb(4)
-              .setNumCpus(5)
-              .setMetadata(ImmutableSet.<Metadata>of())
-              .setExecutorConfig(new ExecutorConfig()
-                  .setName("cmd.exe")
-                  .setData("echo hello world")))
-  );
+          .setTaskConfig(TaskTestUtil.makeConfig(AURORA_JOB_KEY)
+              .newBuilder()
+              .setIsService(false)
+              .setContainer(Container.mesos(new MesosContainer()))));
   static final JobKey QUARTZ_JOB_KEY = Quartz.jobKey(AURORA_JOB_KEY);
 
   private QuartzTestUtil() {

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/filter/SchedulingFilterImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/filter/SchedulingFilterImplTest.java b/src/test/java/org/apache/aurora/scheduler/filter/SchedulingFilterImplTest.java
index 3b1766d..17101f1 100644
--- a/src/test/java/org/apache/aurora/scheduler/filter/SchedulingFilterImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/filter/SchedulingFilterImplTest.java
@@ -27,7 +27,6 @@ import org.apache.aurora.gen.Attribute;
 import org.apache.aurora.gen.Constraint;
 import org.apache.aurora.gen.ExecutorConfig;
 import org.apache.aurora.gen.HostAttributes;
-import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.LimitConstraint;
 import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.TaskConfig;
@@ -35,6 +34,7 @@ import org.apache.aurora.gen.TaskConstraint;
 import org.apache.aurora.gen.ValueConstraint;
 import org.apache.aurora.scheduler.ResourceSlot;
 import org.apache.aurora.scheduler.Resources;
+import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.filter.SchedulingFilter.ResourceRequest;
 import org.apache.aurora.scheduler.filter.SchedulingFilter.UnusedResource;
 import org.apache.aurora.scheduler.filter.SchedulingFilter.Veto;
@@ -44,6 +44,7 @@ import org.apache.aurora.scheduler.mesos.Offers;
 import org.apache.aurora.scheduler.mesos.TaskExecutors;
 import org.apache.aurora.scheduler.storage.entities.IAttribute;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
+import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
 import org.junit.Before;
 import org.junit.Test;
@@ -67,16 +68,8 @@ public class SchedulingFilterImplTest extends EasyMockTest {
   private static final String RACK_ATTRIBUTE = "rack";
   private static final String HOST_ATTRIBUTE = "host";
 
-  private static final String JOB_A = "myJobA";
-  private static final String JOB_B = "myJobB";
-
-  private static final String ROLE_A = "roleA";
-  private static final String USER_A = "userA";
-  private static final Identity OWNER_A = new Identity(ROLE_A, USER_A);
-
-  private static final String ROLE_B = "roleB";
-  private static final String USER_B = "userB";
-  private static final Identity OWNER_B = new Identity(ROLE_B, USER_B);
+  private static final IJobKey JOB_A = JobKeys.from("roleA", "env", "jobA");
+  private static final IJobKey JOB_B = JobKeys.from("roleB", "env", "jobB");
 
   private static final int DEFAULT_CPUS = 4;
   private static final long DEFAULT_RAM = 1000;
@@ -164,22 +157,21 @@ public class SchedulingFilterImplTest extends EasyMockTest {
   public void testDedicatedRole() {
     control.replay();
 
-    IHostAttributes hostA = hostAttributes(HOST_A, dedicated(ROLE_A));
-    checkConstraint(hostA, DEDICATED_ATTRIBUTE, true, ROLE_A);
-    assertVetoes(makeTask(OWNER_B, JOB_B), hostA, Veto.dedicatedHostConstraintMismatch());
+    IHostAttributes hostA = hostAttributes(HOST_A, dedicated(JOB_A.getRole()));
+    checkConstraint(hostA, DEDICATED_ATTRIBUTE, true, JOB_A.getRole());
+    assertVetoes(makeTask(JOB_B), hostA, Veto.dedicatedHostConstraintMismatch());
   }
 
   @Test
   public void testSharedDedicatedHost() {
     control.replay();
 
-    String dedicated1 = "userA/jobA";
-    String dedicated2 = "kestrel/kestrel";
+    String dedicated1 = dedicatedFor(JOB_A);
+    String dedicated2 = dedicatedFor(JOB_B);
     IHostAttributes hostA = hostAttributes(HOST_A, dedicated(dedicated1, dedicated2));
     assertNoVetoes(
         checkConstraint(
-            new Identity().setRole("userA"),
-            "jobA",
+            JOB_A,
             hostA,
             DEDICATED_ATTRIBUTE,
             true,
@@ -187,8 +179,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
         hostA);
     assertNoVetoes(
         checkConstraint(
-            new Identity().setRole("kestrel"),
-            "kestrel",
+            JOB_B,
             hostA,
             DEDICATED_ATTRIBUTE,
             true,
@@ -246,11 +237,11 @@ public class SchedulingFilterImplTest extends EasyMockTest {
     Constraint constraint2 = makeConstraint(DEDICATED_ATTRIBUTE, "xxx");
 
     assertVetoes(
-        makeTask(OWNER_A, JOB_A, constraint1, constraint2),
+        makeTask(JOB_A, constraint1, constraint2),
         hostAttributes(HOST_A, dedicated(HOST_A), host(HOST_A)),
         Veto.constraintMismatch(DEDICATED_ATTRIBUTE));
     assertNoVetoes(
-        makeTask(OWNER_B, JOB_B, constraint1, constraint2),
+        makeTask(JOB_B, constraint1, constraint2),
         hostAttributes(HOST_B, dedicated("xxx"), host(HOST_A)));
   }
 
@@ -263,12 +254,12 @@ public class SchedulingFilterImplTest extends EasyMockTest {
 
     Constraint hostLimit = limitConstraint("host", 1);
     assertVetoes(
-        makeTask(OWNER_A, JOB_A, hostLimit, makeConstraint(DEDICATED_ATTRIBUTE, "xxx")),
+        makeTask(JOB_A, hostLimit, makeConstraint(DEDICATED_ATTRIBUTE, "xxx")),
         hostAttributes(HOST_A, host(HOST_A)),
         Veto.constraintMismatch(DEDICATED_ATTRIBUTE));
     assertVetoes(
-        makeTask(OWNER_B, JOB_A, hostLimit, makeConstraint(DEDICATED_ATTRIBUTE, "xxx")),
-        hostAttributes(HOST_B, dedicated(OWNER_B.getRole() + "/" + JOB_B), host(HOST_B)),
+        makeTask(JOB_A, hostLimit, makeConstraint(DEDICATED_ATTRIBUTE, "xxx")),
+        hostAttributes(HOST_B, dedicated(dedicatedFor(JOB_B)), host(HOST_B)),
         Veto.constraintMismatch(DEDICATED_ATTRIBUTE));
   }
 
@@ -291,6 +282,10 @@ public class SchedulingFilterImplTest extends EasyMockTest {
     return valueAttribute(DEDICATED_ATTRIBUTE, value, values);
   }
 
+  private String dedicatedFor(IJobKey job) {
+    return job.getRole() + "/" + job.getName();
+  }
+
   @Test
   public void testLimitWithinJob() throws Exception {
     control.replay();
@@ -317,39 +312,39 @@ public class SchedulingFilterImplTest extends EasyMockTest {
     IHostAttributes hostA = hostAttributes(HOST_A, host(HOST_A), rack(RACK_A));
     IHostAttributes hostB = hostAttributes(HOST_B, host(HOST_B), rack(RACK_A));
     IHostAttributes hostC = hostAttributes(HOST_C, host(HOST_C), rack(RACK_B));
-    assertNoVetoes(hostLimitTask(OWNER_A, JOB_A, 2), hostA, stateA);
+    assertNoVetoes(hostLimitTask(JOB_A, 2), hostA, stateA);
     assertVetoes(
-        hostLimitTask(OWNER_A, JOB_A, 1),
+        hostLimitTask(JOB_A, 1),
         hostB,
         stateA,
         Veto.unsatisfiedLimit(HOST_ATTRIBUTE));
     assertVetoes(
-        hostLimitTask(OWNER_A, JOB_A, 2),
+        hostLimitTask(JOB_A, 2),
         hostB,
         stateA,
         Veto.unsatisfiedLimit(HOST_ATTRIBUTE));
-    assertNoVetoes(hostLimitTask(OWNER_A, JOB_A, 3), hostB, stateA);
+    assertNoVetoes(hostLimitTask(JOB_A, 3), hostB, stateA);
 
     assertVetoes(
-        rackLimitTask(OWNER_B, JOB_A, 2),
+        rackLimitTask(JOB_A, 2),
         hostB,
         stateB,
         Veto.unsatisfiedLimit(RACK_ATTRIBUTE));
     assertVetoes(
-        rackLimitTask(OWNER_B, JOB_A, 3),
+        rackLimitTask(JOB_A, 3),
         hostB,
         stateB,
         Veto.unsatisfiedLimit(RACK_ATTRIBUTE));
-    assertNoVetoes(rackLimitTask(OWNER_B, JOB_A, 4), hostB, stateB);
+    assertNoVetoes(rackLimitTask(JOB_A, 4), hostB, stateB);
 
-    assertNoVetoes(rackLimitTask(OWNER_B, JOB_A, 1), hostC, stateB);
+    assertNoVetoes(rackLimitTask(JOB_A, 1), hostC, stateB);
 
     assertVetoes(
-        rackLimitTask(OWNER_A, JOB_A, 1),
+        rackLimitTask(JOB_A, 1),
         hostC,
         stateA,
         Veto.unsatisfiedLimit(RACK_ATTRIBUTE));
-    assertNoVetoes(rackLimitTask(OWNER_B, JOB_A, 2), hostC, stateB);
+    assertNoVetoes(rackLimitTask(JOB_A, 2), hostC, stateB);
   }
 
   @Test
@@ -402,7 +397,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
     Constraint jvmConstraint = makeConstraint("jvm", "1.6");
     Constraint zoneConstraint = makeConstraint("zone", "c");
 
-    ITaskConfig task = makeTask(OWNER_A, JOB_A, jvmConstraint, zoneConstraint);
+    ITaskConfig task = makeTask(JOB_A, jvmConstraint, zoneConstraint);
     assertEquals(
         ImmutableSet.of(),
         defaultFilter.filter(
@@ -414,7 +409,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
     Constraint zoneNegated = jvmConstraint.deepCopy();
     zoneNegated.getConstraint().getValue().setNegated(true);
     assertVetoes(
-        makeTask(OWNER_A, JOB_A, jvmNegated, zoneNegated),
+        makeTask(JOB_A, jvmNegated, zoneNegated),
         hostA,
         Veto.constraintMismatch("jvm"));
   }
@@ -476,19 +471,11 @@ public class SchedulingFilterImplTest extends EasyMockTest {
       String value,
       String... vs) {
 
-    return checkConstraint(
-        OWNER_A,
-        JOB_A,
-        hostAttributes,
-        constraintName,
-        expected,
-        value,
-        vs);
+    return checkConstraint(JOB_A, hostAttributes, constraintName, expected, value, vs);
   }
 
   private ITaskConfig checkConstraint(
-      Identity owner,
-      String jobName,
+      IJobKey job,
       IHostAttributes hostAttributes,
       String constraintName,
       boolean expected,
@@ -496,8 +483,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
       String... vs) {
 
     return checkConstraint(
-        owner,
-        jobName,
+        job,
         EMPTY,
         hostAttributes,
         constraintName,
@@ -507,8 +493,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
   }
 
   private ITaskConfig checkConstraint(
-      Identity owner,
-      String jobName,
+      IJobKey job,
       AttributeAggregate aggregate,
       IHostAttributes hostAttributes,
       String constraintName,
@@ -516,7 +501,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
       ValueConstraint value) {
 
     Constraint constraint = new Constraint(constraintName, TaskConstraint.value(value));
-    ITaskConfig task = makeTask(owner, jobName, constraint);
+    ITaskConfig task = makeTask(job, constraint);
     assertEquals(
         expected,
         defaultFilter.filter(
@@ -526,7 +511,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
 
     Constraint negated = constraint.deepCopy();
     negated.getConstraint().getValue().setNegated(!value.isNegated());
-    ITaskConfig negatedTask = makeTask(owner, jobName, negated);
+    ITaskConfig negatedTask = makeTask(job, negated);
     assertEquals(
         !expected,
         defaultFilter.filter(
@@ -598,36 +583,27 @@ public class SchedulingFilterImplTest extends EasyMockTest {
     return new Constraint(name, TaskConstraint.limit(new LimitConstraint(value)));
   }
 
-  private ITaskConfig makeTask(Identity owner, String jobName, Constraint... constraint) {
-    return ITaskConfig.build(makeTask(OWNER_A, JOB_A, DEFAULT_CPUS, DEFAULT_RAM, DEFAULT_DISK)
+  private ITaskConfig makeTask(IJobKey job, Constraint... constraint) {
+    return ITaskConfig.build(makeTask(job, DEFAULT_CPUS, DEFAULT_RAM, DEFAULT_DISK)
         .newBuilder()
-        .setOwner(owner)
-        .setJobName(jobName)
         .setConstraints(Sets.newHashSet(constraint)));
   }
 
-  private ITaskConfig hostLimitTask(Identity owner, String jobName, int maxPerHost) {
-    return makeTask(owner, jobName, limitConstraint(HOST_ATTRIBUTE, maxPerHost));
+  private ITaskConfig hostLimitTask(IJobKey job, int maxPerHost) {
+    return makeTask(job, limitConstraint(HOST_ATTRIBUTE, maxPerHost));
   }
 
   private ITaskConfig hostLimitTask(int maxPerHost) {
-    return hostLimitTask(OWNER_A, JOB_A, maxPerHost);
+    return hostLimitTask(JOB_A, maxPerHost);
   }
 
-  private ITaskConfig rackLimitTask(Identity owner, String jobName, int maxPerRack) {
-    return makeTask(owner, jobName, limitConstraint(RACK_ATTRIBUTE, maxPerRack));
+  private ITaskConfig rackLimitTask(IJobKey job, int maxPerRack) {
+    return makeTask(job, limitConstraint(RACK_ATTRIBUTE, maxPerRack));
   }
 
-  private ITaskConfig makeTask(
-      Identity owner,
-      String jobName,
-      int cpus,
-      long ramMb,
-      long diskMb) {
-
+  private ITaskConfig makeTask(IJobKey job, int cpus, long ramMb, long diskMb) {
     return ITaskConfig.build(new TaskConfig()
-        .setOwner(owner)
-        .setJobName(jobName)
+        .setJob(job.newBuilder())
         .setNumCpus(cpus)
         .setRamMb(ramMb)
         .setDiskMb(diskMb)
@@ -635,7 +611,7 @@ public class SchedulingFilterImplTest extends EasyMockTest {
   }
 
   private ITaskConfig makeTask(int cpus, long ramMb, long diskMb) {
-    return makeTask(OWNER_A, JOB_A, cpus, ramMb, diskMb);
+    return makeTask(JOB_A, cpus, ramMb, diskMb);
   }
 
   private ITaskConfig makeTask() {

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/http/api/ApiBetaTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/http/api/ApiBetaTest.java b/src/test/java/org/apache/aurora/scheduler/http/api/ApiBetaTest.java
index 2b5a82d..f14d405 100644
--- a/src/test/java/org/apache/aurora/scheduler/http/api/ApiBetaTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/http/api/ApiBetaTest.java
@@ -28,15 +28,11 @@ import com.sun.jersey.api.client.config.ClientConfig;
 import com.sun.jersey.api.client.config.DefaultClientConfig;
 
 import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.Constraint;
 import org.apache.aurora.gen.CronCollisionPolicy;
-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.JobSummary;
 import org.apache.aurora.gen.JobSummaryResult;
-import org.apache.aurora.gen.LimitConstraint;
 import org.apache.aurora.gen.Lock;
 import org.apache.aurora.gen.LockKey;
 import org.apache.aurora.gen.Response;
@@ -45,9 +41,8 @@ import org.apache.aurora.gen.RoleSummary;
 import org.apache.aurora.gen.RoleSummaryResult;
 import org.apache.aurora.gen.ScheduleStatusResult;
 import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
-import org.apache.aurora.gen.TaskConstraint;
 import org.apache.aurora.gen.TaskQuery;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.http.AbstractJettyTest;
 import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
 import org.apache.aurora.scheduler.storage.entities.IResponse;
@@ -84,22 +79,7 @@ public class ApiBetaTest extends AbstractJettyTest {
     );
   }
 
-  private static final ITaskConfig TASK_CONFIG = ITaskConfig.build(
-      new TaskConfig()
-          .setOwner(new Identity().setUser("user").setRole("role"))
-          .setEnvironment("test")
-          .setDiskMb(1024)
-          .setRamMb(4096)
-          .setNumCpus(1.0)
-          .setIsService(true)
-          .setJobName("jobName")
-          .setConstraints(ImmutableSet.of(new Constraint()
-              .setName("rack")
-              .setConstraint(TaskConstraint.limit(new LimitConstraint().setLimit(1)))))
-          .setExecutorConfig(
-              new ExecutorConfig()
-                  .setData("executor config data")
-                  .setName("executor name")));
+  private static final ITaskConfig TASK_CONFIG = TaskTestUtil.makeConfig(TaskTestUtil.JOB);
   private static final IJobConfiguration JOB_CONFIG = IJobConfiguration.build(
       new JobConfiguration()
           .setCronCollisionPolicy(CronCollisionPolicy.CANCEL_NEW)

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/mesos/MesosTaskFactoryImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/mesos/MesosTaskFactoryImplTest.java b/src/test/java/org/apache/aurora/scheduler/mesos/MesosTaskFactoryImplTest.java
index 066c6a3..cc2a415 100644
--- a/src/test/java/org/apache/aurora/scheduler/mesos/MesosTaskFactoryImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/mesos/MesosTaskFactoryImplTest.java
@@ -24,14 +24,13 @@ import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.Container;
 import org.apache.aurora.gen.DockerContainer;
 import org.apache.aurora.gen.DockerParameter;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.MesosContainer;
 import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.scheduler.ResourceSlot;
 import org.apache.aurora.scheduler.ResourceType;
 import org.apache.aurora.scheduler.Resources;
 import org.apache.aurora.scheduler.TierManager;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.configuration.executor.ExecutorConfig;
 import org.apache.aurora.scheduler.configuration.executor.ExecutorSettings;
 import org.apache.aurora.scheduler.mesos.MesosTaskFactory.MesosTaskFactoryImpl;
@@ -63,16 +62,10 @@ import static org.junit.Assert.assertTrue;
 
 public class MesosTaskFactoryImplTest extends EasyMockTest {
 
-  private static final ITaskConfig TASK_CONFIG = ITaskConfig.build(new TaskConfig()
-      .setJob(new JobKey("role", "environment", "job-name"))
-      .setOwner(new Identity("role", "user"))
-      .setEnvironment("environment")
-      .setJobName("job-name")
-      .setDiskMb(10)
-      .setRamMb(100)
-      .setNumCpus(5)
-      .setContainer(Container.mesos(new MesosContainer()))
-      .setRequestedPorts(ImmutableSet.of("http")));
+  private static final ITaskConfig TASK_CONFIG = ITaskConfig.build(
+      TaskTestUtil.makeConfig(TaskTestUtil.JOB)
+          .newBuilder()
+          .setContainer(Container.mesos(new MesosContainer())));
   private static final IAssignedTask TASK = IAssignedTask.build(new AssignedTask()
       .setInstanceId(2)
       .setTaskId("task-id")

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java b/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
index e1b5391..ffeee1b 100644
--- a/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
@@ -24,15 +24,11 @@ import org.apache.aurora.common.quantity.Amount;
 import org.apache.aurora.common.quantity.Time;
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
 import org.apache.aurora.common.util.testing.FakeClock;
-import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.ExecutorConfig;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
 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.async.DelayExecutor;
+import org.apache.aurora.scheduler.base.JobKeys;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
 import org.apache.aurora.scheduler.pruning.TaskHistoryPruner.HistoryPrunnerSettings;
@@ -58,7 +54,6 @@ import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 
 public class TaskHistoryPrunerTest extends EasyMockTest {
-  private static final String JOB_A = "job-a";
   private static final String SLAVE_HOST = "HOST_A";
   private static final Amount<Long, Time> ONE_MS = Amount.of(1L, Time.MILLISECONDS);
   private static final Amount<Long, Time> ONE_MINUTE = Amount.of(1L, Time.MINUTES);
@@ -131,7 +126,7 @@ public class TaskHistoryPrunerTest extends EasyMockTest {
 
     clock.advance(ONE_MINUTE);
     IScheduledTask d = makeTask("d", FINISHED);
-    IScheduledTask e = makeTask("job-x", "e", FINISHED);
+    IScheduledTask e = makeTask(JobKeys.from("role", "env", "job-x"), "e", FINISHED);
 
     expectNoImmediatePrune(ImmutableSet.of(a));
     expectOneDelayedPrune(taskATimestamp);
@@ -318,29 +313,19 @@ public class TaskHistoryPrunerTest extends EasyMockTest {
   }
 
   private IScheduledTask makeTask(
-      String job,
+      IJobKey job,
       String taskId,
       ScheduleStatus status) {
 
-    return IScheduledTask.build(new ScheduledTask()
-        .setStatus(status)
-        .setTaskEvents(ImmutableList.of(new TaskEvent(clock.nowMillis(), status)))
-        .setAssignedTask(makeAssignedTask(job, taskId)));
+    ScheduledTask builder = TaskTestUtil.addStateTransition(
+        TaskTestUtil.makeTask(taskId, job), status, clock.nowMillis())
+        .newBuilder();
+    builder.getAssignedTask().setSlaveHost(SLAVE_HOST);
+    return IScheduledTask.build(builder);
   }
 
   private IScheduledTask makeTask(String taskId, ScheduleStatus status) {
-    return makeTask(JOB_A, taskId, status);
+    return makeTask(TaskTestUtil.JOB, taskId, status);
   }
 
-  private AssignedTask makeAssignedTask(String job, String taskId) {
-    return new AssignedTask()
-        .setSlaveHost(SLAVE_HOST)
-        .setTaskId(taskId)
-        .setTask(new TaskConfig()
-            .setJob(new JobKey("role", "staging45", job))
-            .setOwner(new Identity().setRole("role").setUser("user"))
-            .setEnvironment("staging45")
-            .setJobName(job)
-            .setExecutorConfig(new ExecutorConfig("aurora", "config")));
-  }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java b/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
index 920e3e5..aedeab3 100644
--- a/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/quota/QuotaManagerImplTest.java
@@ -20,9 +20,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
-import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.Constraint;
-import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.InstanceTaskConfig;
 import org.apache.aurora.gen.JobConfiguration;
 import org.apache.aurora.gen.JobKey;
@@ -32,13 +30,13 @@ import org.apache.aurora.gen.JobUpdateKey;
 import org.apache.aurora.gen.JobUpdateSummary;
 import org.apache.aurora.gen.Range;
 import org.apache.aurora.gen.ResourceAggregate;
-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.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.quota.QuotaManager.QuotaException;
 import org.apache.aurora.scheduler.quota.QuotaManager.QuotaManagerImpl;
 import org.apache.aurora.scheduler.storage.JobUpdateStore;
@@ -126,13 +124,11 @@ public class QuotaManagerImplTest extends EasyMockTest {
     final String pcRole = "pc-role";
     ScheduledTask ignoredProdTask = prodTask(pcRole, 20, 20, 20).newBuilder();
     ignoredProdTask.getAssignedTask().getTask()
-        .setOwner(new Identity(pcRole, "ignored"))
         .setJob(new JobKey(pcRole, ENV, pcRole));
 
     final String npcRole = "npc-role";
     ScheduledTask ignoredNonProdTask = nonProdTask(npcRole, 20, 20, 20).newBuilder();
     ignoredNonProdTask.getAssignedTask().getTask()
-        .setOwner(new Identity(npcRole, "ignored"))
         .setJob(new JobKey(npcRole, ENV, npcRole));
 
     expectCronJobs(
@@ -976,21 +972,14 @@ public class QuotaManagerImplTest extends EasyMockTest {
       boolean production,
       int instanceId) {
 
-    return IScheduledTask.build(new ScheduledTask()
-        .setStatus(ScheduleStatus.RUNNING)
-        .setAssignedTask(
-            new AssignedTask()
-                .setTaskId(taskId)
-                .setInstanceId(instanceId)
-                .setTask(new TaskConfig()
-                    .setJob(new JobKey(ROLE, ENV, jobName))
-                    .setOwner(new Identity(ROLE, ROLE))
-                    .setEnvironment(ENV)
-                    .setJobName(jobName)
-                    .setNumCpus(cpus)
-                    .setRamMb(ramMb)
-                    .setDiskMb(diskMb)
-                    .setProduction(production))));
+    ScheduledTask builder = TaskTestUtil.makeTask(taskId, JobKeys.from(ROLE, ENV, jobName))
+        .newBuilder();
+    builder.getAssignedTask().setInstanceId(instanceId);
+    builder.getAssignedTask().getTask().setNumCpus(cpus)
+        .setRamMb(ramMb)
+        .setDiskMb(diskMb)
+        .setProduction(production);
+    return IScheduledTask.build(builder);
   }
 
   private IJobConfiguration createJob(IScheduledTask scheduledTask, int instanceCount) {

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/scheduling/RescheduleCalculatorImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/scheduling/RescheduleCalculatorImplTest.java b/src/test/java/org/apache/aurora/scheduler/scheduling/RescheduleCalculatorImplTest.java
index 9d21dcd..2d34729 100644
--- a/src/test/java/org/apache/aurora/scheduler/scheduling/RescheduleCalculatorImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/scheduling/RescheduleCalculatorImplTest.java
@@ -23,12 +23,10 @@ import org.apache.aurora.common.quantity.Amount;
 import org.apache.aurora.common.quantity.Time;
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
 import org.apache.aurora.common.util.BackoffStrategy;
-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.base.TaskTestUtil;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.scheduling.RescheduleCalculator.RescheduleCalculatorImpl;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
@@ -151,14 +149,9 @@ public class RescheduleCalculatorImplTest extends EasyMockTest {
   }
 
   private IScheduledTask makeTask(String taskId) {
-    return IScheduledTask.build(new ScheduledTask()
-        .setAssignedTask(new AssignedTask()
-            .setInstanceId(0)
-            .setTaskId(taskId)
-            .setTask(new TaskConfig()
-                .setJobName("job-" + taskId)
-                .setOwner(new Identity().setRole("role-" + taskId).setUser("user-" + taskId))
-                .setEnvironment("env-" + taskId))));
+    ScheduledTask builder = TaskTestUtil.makeTask(taskId, TaskTestUtil.JOB).newBuilder();
+    builder.unsetAncestorId();
+    return IScheduledTask.build(builder);
   }
 
   private IScheduledTask makeTask(String taskId, ScheduleStatus status) {

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/sla/MetricCalculatorTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/sla/MetricCalculatorTest.java b/src/test/java/org/apache/aurora/scheduler/sla/MetricCalculatorTest.java
index 89162d1..953b65f 100644
--- a/src/test/java/org/apache/aurora/scheduler/sla/MetricCalculatorTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/sla/MetricCalculatorTest.java
@@ -102,7 +102,6 @@ public class MetricCalculatorTest extends EasyMockTest {
         ImmutableSet.of(task1, task2, task3, task4),
         ImmutableSet.of(PROD_METRICS, NON_PROD_METRICS));
 
-    assertEquals(PROD_METRICS.size() + NON_PROD_METRICS.size(), names.getValues().size());
     assertEquals(metricNames, ImmutableSet.copyOf(names.getValues()));
   }
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/sla/SlaTestUtil.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/sla/SlaTestUtil.java b/src/test/java/org/apache/aurora/scheduler/sla/SlaTestUtil.java
index 7fb6278..78f440f 100644
--- a/src/test/java/org/apache/aurora/scheduler/sla/SlaTestUtil.java
+++ b/src/test/java/org/apache/aurora/scheduler/sla/SlaTestUtil.java
@@ -19,13 +19,10 @@ import java.util.Map;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 
-import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.Identity;
-import org.apache.aurora.gen.JobKey;
 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.base.TaskTestUtil;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.entities.ITaskEvent;
 
@@ -41,23 +38,15 @@ final class SlaTestUtil {
 
   static IScheduledTask makeTask(Map<Long, ScheduleStatus> events, int instanceId, boolean isProd) {
     List<ITaskEvent> taskEvents = makeEvents(events);
-    return IScheduledTask.build(new ScheduledTask()
+    ScheduledTask builder = TaskTestUtil.makeTask("task_id", TaskTestUtil.JOB).newBuilder()
         .setStatus(Iterables.getLast(taskEvents).getStatus())
-        .setTaskEvents(ITaskEvent.toBuildersList(taskEvents))
-        .setAssignedTask(new AssignedTask()
-            .setTaskId("task_Id")
-            .setSlaveHost("host")
-            .setInstanceId(instanceId)
-            .setTask(new TaskConfig()
-                .setJob(new JobKey("role", "env", "job"))
-                .setJobName("job")
-                .setIsService(true)
-                .setProduction(isProd)
-                .setEnvironment("env")
-                .setOwner(new Identity("role", "role-user")))));
+        .setTaskEvents(ITaskEvent.toBuildersList(taskEvents));
+    builder.getAssignedTask().setInstanceId(instanceId);
+    builder.getAssignedTask().getTask().setProduction(isProd);
+    return IScheduledTask.build(builder);
   }
 
-  static List<ITaskEvent> makeEvents(Map<Long, ScheduleStatus> events) {
+  private static List<ITaskEvent> makeEvents(Map<Long, ScheduleStatus> events) {
     ImmutableList.Builder<ITaskEvent> taskEvents = ImmutableList.builder();
     for (Map.Entry<Long, ScheduleStatus> entry : events.entrySet()) {
       taskEvents.add(ITaskEvent.build(new TaskEvent(entry.getKey(), entry.getValue())));

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/state/LockManagerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/state/LockManagerImplTest.java b/src/test/java/org/apache/aurora/scheduler/state/LockManagerImplTest.java
index 4a655f1..a146e40 100644
--- a/src/test/java/org/apache/aurora/scheduler/state/LockManagerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/state/LockManagerImplTest.java
@@ -25,7 +25,6 @@ import org.apache.aurora.common.quantity.Amount;
 import org.apache.aurora.common.quantity.Time;
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
 import org.apache.aurora.common.util.testing.FakeClock;
-import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.Lock;
 import org.apache.aurora.gen.LockKey;
 import org.apache.aurora.scheduler.base.JobKeys;
@@ -47,9 +46,8 @@ import static org.junit.Assert.assertEquals;
 
 public class LockManagerImplTest extends EasyMockTest {
   private static final String USER = "jim-user";
-  private static final Identity JIM = new Identity("jim", USER);
   private static final String MY_JOB = "myJob";
-  private static final IJobKey JOB_KEY = JobKeys.from(JIM.getRole(), "devel", MY_JOB);
+  private static final IJobKey JOB_KEY = JobKeys.from("jim", "devel", MY_JOB);
   private static final ILockKey LOCK_KEY = ILockKey.build(LockKey.job(JOB_KEY.newBuilder()));
   private static final UUID TOKEN = UUID.fromString("79d6d790-3212-11e3-aa6e-0800200c9a66");
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java b/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
index 092df8c..94f5ca5 100644
--- a/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
@@ -25,22 +25,21 @@ import com.google.inject.Injector;
 
 import org.apache.aurora.common.stats.StatsProvider;
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
-import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.HostAttributes;
 import org.apache.aurora.gen.HostStatus;
-import org.apache.aurora.gen.Identity;
 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.scheduler.async.AsyncModule.AsyncExecutor;
 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.TaskStateChange;
 import org.apache.aurora.scheduler.events.PubsubEventModule;
 import org.apache.aurora.scheduler.storage.Storage;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
+import org.apache.aurora.scheduler.storage.entities.IHostStatus;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
 import org.apache.aurora.scheduler.testing.FakeStatsProvider;
@@ -83,30 +82,26 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
             bind(StateManager.class).toInstance(stateManager);
             bind(StatsProvider.class).toInstance(new FakeStatsProvider());
             bind(Executor.class).annotatedWith(AsyncExecutor.class)
-                .toInstance(MoreExecutors.sameThreadExecutor());
+                .toInstance(MoreExecutors.directExecutor());
           }
         });
     maintenance = injector.getInstance(MaintenanceController.class);
     eventSink = PubsubTestUtil.startPubsub(injector);
   }
 
-  private static ScheduledTask makeTask(String host, String taskId) {
-    return new ScheduledTask()
-        .setStatus(RUNNING)
-        .setAssignedTask(
-            new AssignedTask()
-                .setSlaveHost(host)
-                .setTaskId(taskId)
-                .setTask(
-                    new TaskConfig()
-                        .setJobName("jobName")
-                        .setOwner(new Identity().setRole("role").setUser("role"))));
+  private static IScheduledTask makeTask(String host, String taskId) {
+    ScheduledTask builder = TaskTestUtil.addStateTransition(
+        TaskTestUtil.makeTask(taskId, TaskTestUtil.JOB),
+        RUNNING,
+        1000).newBuilder();
+    builder.getAssignedTask().setSlaveHost(host);
+    return IScheduledTask.build(builder);
   }
 
   @Test
   public void testMaintenanceCycle() {
-    ScheduledTask task1 = makeTask(HOST_A, "taskA");
-    ScheduledTask task2 = makeTask(HOST_A, "taskB");
+    IScheduledTask task1 = makeTask(HOST_A, "taskA");
+    IScheduledTask task2 = makeTask(HOST_A, "taskB");
 
     expectMaintenanceModeChange(HOST_A, SCHEDULED);
     expectFetchTasksByHost(HOST_A, ImmutableSet.of(task1, task2));
@@ -132,9 +127,11 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
     assertStatus(HOST_A, DRAINING, maintenance.drain(A));
     assertStatus(HOST_A, DRAINING, maintenance.getStatus(A));
     eventSink.post(
-        TaskStateChange.transition(IScheduledTask.build(task1.setStatus(KILLED)), RUNNING));
+        TaskStateChange.transition(
+            IScheduledTask.build(task1.newBuilder().setStatus(KILLED)), RUNNING));
     eventSink.post(
-        TaskStateChange.transition(IScheduledTask.build(task2.setStatus(KILLED)), RUNNING));
+        TaskStateChange.transition(
+            IScheduledTask.build(task2.newBuilder().setStatus(KILLED)), RUNNING));
     assertStatus(HOST_A, NONE, maintenance.endMaintenance(A));
   }
 
@@ -178,7 +175,7 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
     // Make sure a later transition on the host does not cause any ill effects that could surface
     // from stale internal state.
     eventSink.post(TaskStateChange.transition(
-        IScheduledTask.build(makeTask(HOST_A, "taskA").setStatus(KILLED)), RUNNING));
+        IScheduledTask.build(makeTask(HOST_A, "taskA").newBuilder().setStatus(KILLED)), RUNNING));
   }
 
   @Test
@@ -193,7 +190,7 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
     assertEquals(NONE, maintenance.getMode("unknown"));
   }
 
-  private void expectTaskDraining(ScheduledTask task) {
+  private void expectTaskDraining(IScheduledTask task) {
     expect(stateManager.changeState(
         storageUtil.mutableStoreProvider,
         Tasks.id(task),
@@ -203,9 +200,8 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
         .andReturn(StateChangeResult.SUCCESS);
   }
 
-  private void expectFetchTasksByHost(String hostName, ImmutableSet<ScheduledTask> tasks) {
-    expect(storageUtil.taskStore.fetchTasks(Query.slaveScoped(hostName).active()))
-        .andReturn(IScheduledTask.setFromBuilders(tasks));
+  private void expectFetchTasksByHost(String hostName, Set<IScheduledTask> tasks) {
+    expect(storageUtil.taskStore.fetchTasks(Query.slaveScoped(hostName).active())).andReturn(tasks);
   }
 
   private void expectMaintenanceModeChange(String hostName, MaintenanceMode mode) {
@@ -217,7 +213,7 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
     expect(storageUtil.attributeStore.saveHostAttributes(updated)).andReturn(true);
   }
 
-  private void assertStatus(String host, MaintenanceMode mode, Set<HostStatus> statuses) {
-    assertEquals(ImmutableSet.of(new HostStatus(host, mode)), statuses);
+  private void assertStatus(String host, MaintenanceMode mode, Set<IHostStatus> statuses) {
+    assertEquals(ImmutableSet.of(IHostStatus.build(new HostStatus(host, mode))), statuses);
   }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/state/TaskStateMachineTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/state/TaskStateMachineTest.java b/src/test/java/org/apache/aurora/scheduler/state/TaskStateMachineTest.java
index d0a6cd6..c519323 100644
--- a/src/test/java/org/apache/aurora/scheduler/state/TaskStateMachineTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/state/TaskStateMachineTest.java
@@ -25,11 +25,9 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 
-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.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.state.SideEffect.Action;
 import org.apache.aurora.scheduler.state.TaskStateMachine.TaskState;
@@ -315,16 +313,12 @@ public class TaskStateMachineTest {
   }
 
   private static ScheduledTask makeTask(boolean service) {
-    return new ScheduledTask()
-        .setStatus(INIT.getStatus().get())
-        .setAssignedTask(
-            new AssignedTask()
-                .setTaskId("test")
-                .setTask(
-                    new TaskConfig()
-                        .setOwner(new Identity().setRole("roleA"))
-                        .setJobName("jobA")
-                        .setIsService(service)));
+    ScheduledTask builder = TaskTestUtil.makeTask("test", TaskTestUtil.JOB).newBuilder();
+    builder.setStatus(INIT.getStatus().get());
+    builder.getAssignedTask().getTask()
+        .setMaxTaskFailures(0)
+        .setIsService(service);
+    return builder;
   }
 
   private static final TransitionResult SAVE = new TransitionResult(

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
index 58d57b9..b9ce2dd 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
@@ -54,6 +54,7 @@ import org.apache.aurora.gen.JobUpdateSummary;
 import org.apache.aurora.gen.LimitConstraint;
 import org.apache.aurora.gen.ListBackupsResult;
 import org.apache.aurora.gen.LockKey;
+import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.MesosContainer;
 import org.apache.aurora.gen.PulseJobUpdateResult;
 import org.apache.aurora.gen.QueryRecoveryResult;
@@ -91,6 +92,7 @@ import org.apache.aurora.scheduler.state.UUIDGenerator;
 import org.apache.aurora.scheduler.storage.Storage.StorageException;
 import org.apache.aurora.scheduler.storage.backup.Recovery;
 import org.apache.aurora.scheduler.storage.backup.StorageBackup;
+import org.apache.aurora.scheduler.storage.entities.IHostStatus;
 import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
 import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
 import org.apache.aurora.scheduler.storage.entities.IJobKey;
@@ -1233,13 +1235,17 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task), null));
   }
 
+  private static Set<IHostStatus> status(String host, MaintenanceMode mode) {
+    return ImmutableSet.of(IHostStatus.build(new HostStatus(host, mode)));
+  }
+
   @Test
   public void testHostMaintenance() throws Exception {
     Set<String> hostnames = ImmutableSet.of("a");
-    Set<HostStatus> none = ImmutableSet.of(new HostStatus("a", NONE));
-    Set<HostStatus> scheduled = ImmutableSet.of(new HostStatus("a", SCHEDULED));
-    Set<HostStatus> draining = ImmutableSet.of(new HostStatus("a", DRAINING));
-    Set<HostStatus> drained = ImmutableSet.of(new HostStatus("a", DRAINING));
+    Set<IHostStatus> none = status("a", NONE);
+    Set<IHostStatus> scheduled = status("a", SCHEDULED);
+    Set<IHostStatus> draining = status("a", DRAINING);
+    Set<IHostStatus> drained = status("a", DRAINING);
     expect(maintenance.getStatus(hostnames)).andReturn(none);
     expect(maintenance.startMaintenance(hostnames)).andReturn(scheduled);
     expect(maintenance.drain(hostnames)).andReturn(draining);
@@ -1252,26 +1258,26 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     Hosts hosts = new Hosts(hostnames);
 
     assertEquals(
-        none,
+        IHostStatus.toBuildersSet(none),
         thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
             .getStatuses());
     assertEquals(
-        scheduled,
+        IHostStatus.toBuildersSet(scheduled),
         thrift.startMaintenance(hosts).getResult().getStartMaintenanceResult()
             .getStatuses());
     assertEquals(
-        draining,
+        IHostStatus.toBuildersSet(draining),
         thrift.drainHosts(hosts).getResult().getDrainHostsResult().getStatuses());
     assertEquals(
-        draining,
+        IHostStatus.toBuildersSet(draining),
         thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
             .getStatuses());
     assertEquals(
-        drained,
+        IHostStatus.toBuildersSet(drained),
         thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
             .getStatuses());
     assertEquals(
-        none,
+        IHostStatus.toBuildersSet(none),
         thrift.endMaintenance(hosts).getResult().getEndMaintenanceResult().getStatuses());
   }
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/5684c64c/src/test/java/org/apache/aurora/scheduler/updater/JobUpdateEventSubscriberTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/updater/JobUpdateEventSubscriberTest.java b/src/test/java/org/apache/aurora/scheduler/updater/JobUpdateEventSubscriberTest.java
index fffbc36..36f2c65 100644
--- a/src/test/java/org/apache/aurora/scheduler/updater/JobUpdateEventSubscriberTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/updater/JobUpdateEventSubscriberTest.java
@@ -18,15 +18,10 @@ import com.google.common.eventbus.EventBus;
 import com.google.common.util.concurrent.Service;
 
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
-import org.apache.aurora.gen.AssignedTask;
-import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.InstanceKey;
 import org.apache.aurora.gen.ScheduleStatus;
-import org.apache.aurora.gen.ScheduledTask;
-import org.apache.aurora.gen.TaskConfig;
-import org.apache.aurora.scheduler.base.JobKeys;
+import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
-import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.junit.Before;
 import org.junit.Test;
@@ -37,22 +32,10 @@ import static org.easymock.EasyMock.expectLastCall;
 
 public class JobUpdateEventSubscriberTest extends EasyMockTest {
 
-  private static final IJobKey JOB = JobKeys.from("role", "env", "name");
-
-  private static final IScheduledTask TASK = IScheduledTask.build(
-      new ScheduledTask()
-          .setStatus(ScheduleStatus.PENDING)
-          .setAssignedTask(
-              new AssignedTask()
-                  .setInstanceId(5)
-                  .setTask(new TaskConfig()
-                      .setJob(JOB.newBuilder())
-                      .setOwner(new Identity().setRole(JOB.getRole()))
-                      .setEnvironment(JOB.getEnvironment())
-                      .setJobName(JOB.getName()))));
+  private static final IScheduledTask TASK = TaskTestUtil.makeTask("id", TaskTestUtil.JOB);
   private static final IInstanceKey INSTANCE_A = IInstanceKey.build(
       new InstanceKey()
-          .setJobKey(JOB.newBuilder())
+          .setJobKey(TaskTestUtil.JOB.newBuilder())
           .setInstanceId(TASK.getAssignedTask().getInstanceId()));
 
   private JobUpdateController updater;