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 <[email protected]> Authored: Fri Jan 22 17:14:28 2016 -0800 Committer: Bill Farner <[email protected]> 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;
