aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From wfar...@apache.org
Subject [1/2] incubator-aurora git commit: Refactor existing write APIs for job updates to use IJobUpdateKey.
Date Sat, 21 Feb 2015 01:44:41 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master e5de618cc -> ace02d131


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/ace02d13/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java b/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
index 0a5cc51..7e5d0ed 100644
--- a/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
@@ -45,6 +45,7 @@ import org.apache.aurora.gen.JobUpdate;
 import org.apache.aurora.gen.JobUpdateAction;
 import org.apache.aurora.gen.JobUpdateEvent;
 import org.apache.aurora.gen.JobUpdateInstructions;
+import org.apache.aurora.gen.JobUpdateKey;
 import org.apache.aurora.gen.JobUpdateSettings;
 import org.apache.aurora.gen.JobUpdateStatus;
 import org.apache.aurora.gen.JobUpdateSummary;
@@ -97,6 +98,8 @@ import org.apache.aurora.scheduler.storage.entities.IJobInstanceUpdateEvent;
 import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdateEvent;
+import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
+import org.apache.aurora.scheduler.storage.entities.IJobUpdateSummary;
 import org.apache.aurora.scheduler.storage.entities.ILock;
 import org.apache.aurora.scheduler.storage.entities.ILockKey;
 import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
@@ -106,6 +109,7 @@ import org.apache.aurora.scheduler.storage.log.LogStorage.SchedulingService;
 import org.apache.aurora.scheduler.storage.log.testing.LogOpMatcher;
 import org.apache.aurora.scheduler.storage.log.testing.LogOpMatcher.StreamMatcher;
 import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
+import org.apache.aurora.scheduler.updater.Updates;
 import org.easymock.Capture;
 import org.easymock.IAnswer;
 import org.junit.Before;
@@ -124,7 +128,8 @@ public class LogStorageTest extends EasyMockTest {
 
   private static final Amount<Long, Time> SNAPSHOT_INTERVAL = Amount.of(1L, Time.MINUTES);
   private static final IJobKey JOB_KEY = JobKeys.from("role", "env", "name");
-  private static final String UPDATE_ID = "testUpdateId";
+  private static final IJobUpdateKey UPDATE_ID =
+      IJobUpdateKey.build(new JobUpdateKey(JOB_KEY.newBuilder(), "testUpdateId"));
   private static final long NOW = 42L;
 
   private LogStorage logStorage;
@@ -347,27 +352,48 @@ public class LogStorageTest extends EasyMockTest {
     builder.add(createTransaction(Op.removeLock(removeLock)));
     storageUtil.lockStore.removeLock(ILockKey.build(removeLock.getLockKey()));
 
-    SaveJobUpdate saveUpdate = new SaveJobUpdate(new JobUpdate(), "token");
+    JobUpdate update = new JobUpdate()
+        .setSummary(new JobUpdateSummary()
+            .setJobKey(UPDATE_ID.getJob().newBuilder())
+            .setUpdateId(UPDATE_ID.getId()));
+    SaveJobUpdate saveUpdate = new SaveJobUpdate(update, "token");
     builder.add(createTransaction(Op.saveJobUpdate(saveUpdate)));
     storageUtil.jobUpdateStore.saveJobUpdate(
         IJobUpdate.build(saveUpdate.getJobUpdate()),
         Optional.of(saveUpdate.getLockToken()));
 
-    SaveJobUpdateEvent saveUpdateEvent = new SaveJobUpdateEvent(new JobUpdateEvent(), "update");
+    IJobUpdateKey key = Updates.getKey(IJobUpdateSummary.build(update.getSummary()));
+    SaveJobUpdateEvent saveUpdateEvent =
+        new SaveJobUpdateEvent(new JobUpdateEvent(), "update", key.newBuilder());
     builder.add(createTransaction(Op.saveJobUpdateEvent(saveUpdateEvent)));
     storageUtil.jobUpdateStore.saveJobUpdateEvent(
-        IJobUpdateEvent.build(saveUpdateEvent.getEvent()),
-        saveUpdateEvent.getUpdateId());
+        key,
+        IJobUpdateEvent.build(saveUpdateEvent.getEvent()));
 
     SaveJobInstanceUpdateEvent saveInstanceEvent =
-        new SaveJobInstanceUpdateEvent(new JobInstanceUpdateEvent(), "update");
+        new SaveJobInstanceUpdateEvent(new JobInstanceUpdateEvent(), "update", key.newBuilder());
     builder.add(createTransaction(Op.saveJobInstanceUpdateEvent(saveInstanceEvent)));
     storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(
-        IJobInstanceUpdateEvent.build(saveInstanceEvent.getEvent()),
-        saveInstanceEvent.getUpdateId());
+        key,
+        IJobInstanceUpdateEvent.build(saveInstanceEvent.getEvent()));
+
+    // Backwards compatibility as part of AURORA-1093.  Exercises behavior to drop job
+    // update-related records that cannot be found based on their update ID string alone.
+    SaveJobInstanceUpdateEvent unknownSaveInstanceEvent =
+        new SaveJobInstanceUpdateEvent(new JobInstanceUpdateEvent(), "update5", null);
+    builder.add(createTransaction(Op.saveJobInstanceUpdateEvent(unknownSaveInstanceEvent)));
+    expect(storageUtil.jobUpdateStore.fetchUpdateKey("update5"))
+        .andReturn(Optional.<IJobUpdateKey>absent());
+
+    SaveJobUpdateEvent unknownSaveUpdateEvent =
+        new SaveJobUpdateEvent(new JobUpdateEvent(), "update6", null);
+    builder.add(createTransaction(Op.saveJobUpdateEvent(unknownSaveUpdateEvent)));
+    expect(storageUtil.jobUpdateStore.fetchUpdateKey("update6"))
+        .andReturn(Optional.<IJobUpdateKey>absent());
 
     builder.add(createTransaction(Op.pruneJobUpdateHistory(new PruneJobUpdateHistory(5, 10L))));
-    expect(storageUtil.jobUpdateStore.pruneHistory(5, 10L)).andReturn(ImmutableSet.of("id2"));
+    expect(storageUtil.jobUpdateStore.pruneHistory(5, 10L))
+        .andReturn(ImmutableSet.of(IJobUpdateKey.build(UPDATE_ID.newBuilder())));
 
     // NOOP LogEntry
     builder.add(LogEntry.noop(true));
@@ -928,15 +954,16 @@ public class LogStorageTest extends EasyMockTest {
       @Override
       protected void setupExpectations() throws Exception {
         storageUtil.expectWrite();
-        storageUtil.jobUpdateStore.saveJobUpdateEvent(event, UPDATE_ID);
+        storageUtil.jobUpdateStore.saveJobUpdateEvent(UPDATE_ID, event);
         streamMatcher.expectTransaction(Op.saveJobUpdateEvent(new SaveJobUpdateEvent(
             event.newBuilder(),
-            UPDATE_ID))).andReturn(position);
+            UPDATE_ID.getId(),
+            UPDATE_ID.newBuilder()))).andReturn(position);
       }
 
       @Override
       protected void performMutations(MutableStoreProvider storeProvider) {
-        storeProvider.getJobUpdateStore().saveJobUpdateEvent(event, UPDATE_ID);
+        storeProvider.getJobUpdateStore().saveJobUpdateEvent(UPDATE_ID, event);
       }
     }.run();
   }
@@ -952,14 +979,18 @@ public class LogStorageTest extends EasyMockTest {
       @Override
       protected void setupExpectations() throws Exception {
         storageUtil.expectWrite();
-        storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(event, UPDATE_ID);
+        storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(UPDATE_ID, event);
         streamMatcher.expectTransaction(Op.saveJobInstanceUpdateEvent(
-            new SaveJobInstanceUpdateEvent(event.newBuilder(), UPDATE_ID))).andReturn(position);
+            new SaveJobInstanceUpdateEvent(
+                event.newBuilder(),
+                UPDATE_ID.getId(),
+                UPDATE_ID.newBuilder())))
+            .andReturn(position);
       }
 
       @Override
       protected void performMutations(MutableStoreProvider storeProvider) {
-        storeProvider.getJobUpdateStore().saveJobInstanceUpdateEvent(event, UPDATE_ID);
+        storeProvider.getJobUpdateStore().saveJobInstanceUpdateEvent(UPDATE_ID, event);
       }
     }.run();
   }
@@ -976,7 +1007,8 @@ public class LogStorageTest extends EasyMockTest {
         storageUtil.expectWrite();
         expect(storageUtil.jobUpdateStore.pruneHistory(
             pruneHistory.getPerJobRetainCount(),
-            pruneHistory.getHistoryPruneThresholdMs())).andReturn(ImmutableSet.of("id1"));
+            pruneHistory.getHistoryPruneThresholdMs()))
+            .andReturn(ImmutableSet.of(UPDATE_ID));
 
         streamMatcher.expectTransaction(Op.pruneJobUpdateHistory(pruneHistory)).andReturn(position);
       }

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/ace02d13/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplTest.java
b/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplTest.java
index 8fc3cb8..152e704 100644
--- a/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/storage/log/SnapshotStoreImplTest.java
@@ -31,6 +31,7 @@ import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.JobUpdate;
 import org.apache.aurora.gen.JobUpdateDetails;
 import org.apache.aurora.gen.JobUpdateEvent;
+import org.apache.aurora.gen.JobUpdateKey;
 import org.apache.aurora.gen.JobUpdateStatus;
 import org.apache.aurora.gen.JobUpdateSummary;
 import org.apache.aurora.gen.Lock;
@@ -49,6 +50,7 @@ import org.apache.aurora.scheduler.storage.SnapshotStore;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
 import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdateDetails;
+import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
 import org.apache.aurora.scheduler.storage.entities.ILock;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
@@ -74,6 +76,11 @@ public class SnapshotStoreImplTest extends EasyMockTest {
     snapshotStore = new SnapshotStoreImpl(clock, storageUtil.storage);
   }
 
+  private static IJobUpdateKey makeKey(String id) {
+    return IJobUpdateKey.build(
+        new JobUpdateKey(JobKeys.from("role", "env", "job").newBuilder(), id));
+  }
+
   @Test
   public void testCreateAndRestoreNewSnapshot() {
     ImmutableSet<IScheduledTask> tasks = ImmutableSet.of(
@@ -100,15 +107,21 @@ public class SnapshotStoreImplTest extends EasyMockTest {
     SchedulerMetadata metadata = new SchedulerMetadata()
         .setFrameworkId(frameworkId)
         .setVersion(CURRENT_API_VERSION);
-    final String updateId1 = "updateId1";
-    final String updateId2 = "updateId2";
+    IJobUpdateKey updateId1 =  makeKey("updateId1");
+    IJobUpdateKey updateId2 = makeKey("updateId2");
     IJobUpdateDetails updateDetails1 = IJobUpdateDetails.build(new JobUpdateDetails()
-        .setUpdate(new JobUpdate().setSummary(new JobUpdateSummary().setUpdateId(updateId1)))
+        .setUpdate(new JobUpdate().setSummary(
+            new JobUpdateSummary()
+                .setUpdateId(updateId1.getId())
+                .setJobKey(updateId1.getJob().newBuilder())))
         .setUpdateEvents(ImmutableList.of(new JobUpdateEvent().setStatus(JobUpdateStatus.ERROR)))
         .setInstanceEvents(ImmutableList.of(new JobInstanceUpdateEvent().setTimestampMs(123L))));
 
     IJobUpdateDetails updateDetails2 = IJobUpdateDetails.build(new JobUpdateDetails()
-        .setUpdate(new JobUpdate().setSummary(new JobUpdateSummary().setUpdateId(updateId2))));
+        .setUpdate(new JobUpdate().setSummary(
+            new JobUpdateSummary()
+                .setUpdateId(updateId2.getId())
+                .setJobKey(updateId2.getJob().newBuilder()))));
 
     storageUtil.expectOperations();
     expect(storageUtil.taskStore.fetchTasks(Query.unscoped())).andReturn(tasks);
@@ -139,11 +152,11 @@ public class SnapshotStoreImplTest extends EasyMockTest {
     storageUtil.jobUpdateStore.saveJobUpdate(
         updateDetails1.getUpdate(), Optional.fromNullable(lockToken));
     storageUtil.jobUpdateStore.saveJobUpdateEvent(
-        Iterables.getOnlyElement(updateDetails1.getUpdateEvents()),
-        updateId1);
+        updateId1,
+        Iterables.getOnlyElement(updateDetails1.getUpdateEvents()));
     storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(
-        Iterables.getOnlyElement(updateDetails1.getInstanceEvents()),
-        updateId1);
+        updateId1,
+        Iterables.getOnlyElement(updateDetails1.getInstanceEvents()));
     storageUtil.jobUpdateStore.saveJobUpdate(
         updateDetails2.getUpdate(), Optional.<String>absent());
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/ace02d13/src/test/java/org/apache/aurora/scheduler/storage/log/WriteAheadStorageTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/log/WriteAheadStorageTest.java
b/src/test/java/org/apache/aurora/scheduler/storage/log/WriteAheadStorageTest.java
index 9393ad7..36809ac 100644
--- a/src/test/java/org/apache/aurora/scheduler/storage/log/WriteAheadStorageTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/storage/log/WriteAheadStorageTest.java
@@ -24,12 +24,14 @@ import com.twitter.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.JobUpdateKey;
 import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.storage.Op;
 import org.apache.aurora.gen.storage.PruneJobUpdateHistory;
 import org.apache.aurora.gen.storage.SaveHostAttributes;
 import org.apache.aurora.gen.storage.SaveTasks;
+import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.events.EventSink;
 import org.apache.aurora.scheduler.events.PubsubEvent;
@@ -41,6 +43,7 @@ import org.apache.aurora.scheduler.storage.QuotaStore;
 import org.apache.aurora.scheduler.storage.SchedulerStore;
 import org.apache.aurora.scheduler.storage.TaskStore;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
+import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.easymock.EasyMock;
 import org.junit.Before;
@@ -98,7 +101,10 @@ public class WriteAheadStorageTest extends EasyMockTest {
 
   @Test
   public void testPruneHistory() {
-    Set<String> pruned = ImmutableSet.of("a", "b");
+    Set<IJobUpdateKey> pruned = ImmutableSet.of(
+        IJobUpdateKey.build(new JobUpdateKey(JobKeys.from("role", "env", "job").newBuilder(),
"a")),
+        IJobUpdateKey.build(
+            new JobUpdateKey(JobKeys.from("role", "env", "job").newBuilder(), "b")));
     expect(jobUpdateStore.pruneHistory(1, 1)).andReturn(pruned);
     expectOp(Op.pruneJobUpdateHistory(new PruneJobUpdateHistory(1, 1)));
 
@@ -109,7 +115,7 @@ public class WriteAheadStorageTest extends EasyMockTest {
 
   @Test
   public void testNoopPruneHistory() {
-    expect(jobUpdateStore.pruneHistory(1, 1)).andReturn(ImmutableSet.<String>of());
+    expect(jobUpdateStore.pruneHistory(1, 1)).andReturn(ImmutableSet.<IJobUpdateKey>of());
 
     control.replay();
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/ace02d13/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 06c8faa..1514775 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
@@ -207,10 +207,9 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
   private static final JobConfiguration CRON_JOB = makeJob().setCronSchedule("* * * * *");
   private static final Lock DEFAULT_LOCK = null;
   private static final String TASK_ID = "task_id";
-  private static final String UPDATE_ID = "82d6d790-3212-11e3-aa6e-0800200c9a74";
-  private static final UUID UU_ID = UUID.fromString(UPDATE_ID);
+  private static final UUID UU_ID = UUID.fromString("82d6d790-3212-11e3-aa6e-0800200c9a74");
   private static final IJobUpdateKey UPDATE_KEY =
-      IJobUpdateKey.build(new JobUpdateKey(JOB_KEY.newBuilder(), UPDATE_ID));
+      IJobUpdateKey.build(new JobUpdateKey(JOB_KEY.newBuilder(), UU_ID.toString()));
 
   private static final IResourceAggregate QUOTA =
       IResourceAggregate.build(new ResourceAggregate(10.0, 1024, 2048));
@@ -2621,7 +2620,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
     request.getTaskConfig().unsetJob();
 
     Response response = assertOkResponse(thrift.startJobUpdate(request, SESSION));
-    assertEquals(UPDATE_ID, response.getResult().getStartJobUpdateResult().getUpdateId());
+    assertEquals(UPDATE_KEY.getId(), response.getResult().getStartJobUpdateResult().getUpdateId());
   }
 
   @Test
@@ -2661,7 +2660,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     Response response =
         assertOkResponse(thrift.startJobUpdate(buildJobUpdateRequest(update), SESSION));
-    assertEquals(UPDATE_ID, response.getResult().getStartJobUpdateResult().getUpdateId());
+    assertEquals(UPDATE_KEY.getId(), response.getResult().getStartJobUpdateResult().getUpdateId());
   }
 
   @Test
@@ -3089,7 +3088,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
   @Test
   public void testPulseJobUpdatePulsedAsCoordinator() throws Exception {
     expectAuth(UPDATE_COORDINATOR, true);
-    expect(jobUpdateController.pulse(UPDATE_ID)).andReturn(JobUpdatePulseStatus.OK);
+    expect(jobUpdateController.pulse(UPDATE_KEY)).andReturn(JobUpdatePulseStatus.OK);
 
     control.replay();
 
@@ -3102,7 +3101,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
   public void testPulseJobUpdatePulsedAsUser() throws Exception {
     expectAuth(UPDATE_COORDINATOR, false);
     expectAuth(ROLE, true);
-    expect(jobUpdateController.pulse(UPDATE_ID)).andReturn(JobUpdatePulseStatus.OK);
+    expect(jobUpdateController.pulse(UPDATE_KEY)).andReturn(JobUpdatePulseStatus.OK);
 
     control.replay();
 
@@ -3114,7 +3113,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
   @Test
   public void testPulseJobUpdateFails() throws Exception {
     expectAuth(UPDATE_COORDINATOR, true);
-    expect(jobUpdateController.pulse(UPDATE_ID)).andThrow(new UpdateStateException("failure"));
+    expect(jobUpdateController.pulse(UPDATE_KEY)).andThrow(new UpdateStateException("failure"));
 
     control.replay();
 
@@ -3296,8 +3295,8 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
 
     return IJobUpdate.build(new JobUpdate()
         .setSummary(new JobUpdateSummary()
-            .setJobKey(JOB_KEY.newBuilder())
-            .setUpdateId(UPDATE_ID)
+            .setJobKey(UPDATE_KEY.getJob().newBuilder())
+            .setUpdateId(UPDATE_KEY.getId())
             .setUser(ROLE_IDENTITY.getUser()))
         .setInstructions(new JobUpdateInstructions()
             .setSettings(buildJobUpdateSettings())

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/ace02d13/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java b/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
index a15c409..2088a8f 100644
--- a/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
+++ b/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
@@ -51,6 +51,7 @@ import org.apache.aurora.gen.JobUpdate;
 import org.apache.aurora.gen.JobUpdateAction;
 import org.apache.aurora.gen.JobUpdateEvent;
 import org.apache.aurora.gen.JobUpdateInstructions;
+import org.apache.aurora.gen.JobUpdateKey;
 import org.apache.aurora.gen.JobUpdatePulseStatus;
 import org.apache.aurora.gen.JobUpdateSettings;
 import org.apache.aurora.gen.JobUpdateStatus;
@@ -87,6 +88,7 @@ import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdateDetails;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdateEvent;
+import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
 import org.apache.aurora.scheduler.storage.entities.IJobUpdateSummary;
 import org.apache.aurora.scheduler.storage.entities.ILock;
 import org.apache.aurora.scheduler.storage.entities.ILockKey;
@@ -132,9 +134,10 @@ import static org.junit.Assert.fail;
 
 public class JobUpdaterIT extends EasyMockTest {
 
-  private static final String UPDATE_ID = "update_id";
   private static final String USER = "user";
   private static final IJobKey JOB = JobKeys.from("role", "env", "job1");
+  private static final IJobUpdateKey UPDATE_ID =
+      IJobUpdateKey.build(new JobUpdateKey(JOB.newBuilder(), "update_id"));
   private static final Amount<Long, Time> RUNNING_TIMEOUT = Amount.of(1000L, Time.MILLISECONDS);
   private static final Amount<Long, Time> WATCH_TIMEOUT = Amount.of(2000L, Time.MILLISECONDS);
   private static final Amount<Long, Time> FLAPPING_THRESHOLD = Amount.of(1L, Time.MILLISECONDS);
@@ -277,7 +280,7 @@ public class JobUpdaterIT extends EasyMockTest {
     IJobUpdateDetails details = storage.read(new Work.Quiet<IJobUpdateDetails>() {
       @Override
       public IJobUpdateDetails apply(StoreProvider storeProvider) {
-        return storeProvider.getJobUpdateStore().fetchJobUpdateDetails(UPDATE_ID).get();
+        return storeProvider.getJobUpdateStore().fetchJobUpdateDetails(UPDATE_ID.getId()).get();
       }
     });
     Iterable<IJobInstanceUpdateEvent> orderedEvents =
@@ -624,7 +627,9 @@ public class JobUpdaterIT extends EasyMockTest {
   public void testPulseInvalidUpdateId() throws Exception {
     control.replay();
 
-    assertEquals(JobUpdatePulseStatus.FINISHED, updater.pulse("invalid"));
+    assertEquals(
+        JobUpdatePulseStatus.FINISHED,
+        updater.pulse(IJobUpdateKey.build(new JobUpdateKey(JOB.newBuilder(), "invalid"))));
   }
 
   @Test
@@ -1012,11 +1017,11 @@ public class JobUpdaterIT extends EasyMockTest {
 
     store.saveJobUpdate(update, Optional.of(lock.getToken()));
     store.saveJobUpdateEvent(
+        Updates.getKey(update.getSummary()),
         IJobUpdateEvent.build(
             new JobUpdateEvent()
                 .setStatus(status)
-                .setTimestampMs(clock.nowMillis())),
-        update.getSummary().getUpdateId());
+                .setTimestampMs(clock.nowMillis())));
     return lock;
   }
 
@@ -1287,7 +1292,7 @@ public class JobUpdaterIT extends EasyMockTest {
     return IJobUpdateSummary.build(new JobUpdateSummary()
         .setUser("user")
         .setJobKey(JOB.newBuilder())
-        .setUpdateId(UPDATE_ID));
+        .setUpdateId(UPDATE_ID.getId()));
   }
 
   private static IJobUpdate makeJobUpdate(IInstanceTaskConfig... configs) {


Mime
View raw message