content: /**
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.aurora.scheduler.updater;
  
  import java.util.Set;
  
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.aurora.gen.AssignedTask;
  import org.apache.aurora.gen.ExecutorConfig;
  import org.apache.aurora.gen.Identity;
  import org.apache.aurora.gen.Range;
  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.Query;
  import org.apache.aurora.scheduler.storage.TaskStore;
  import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IRange;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
  import org.easymock.IExpectationSetters;
  import org.junit.Before;
  import org.junit.Test;
  
  import static org.easymock.EasyMock.expect;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertNotEquals;
  
  public class JobDiffTest extends EasyMockTest {
  
    private static final IJobKey JOB = JobKeys.from("role", "env", "job");
    private static final Set<IRange> NO_SCOPE = ImmutableSet.of();
    private static final Set<IRange> CANARY_SCOPE = ImmutableSet.of(IRange.build(new Range(0, 0)));
  
    private TaskStore store;
  
    @Before
    public void setUp() {
      store = createMock(TaskStore.class);
    }
  
    @Test
    public void testNoDiff() {
      ITaskConfig task = makeTask("job", "echo");
  
      expectFetch(instance(task, 0), instance(task, 1)).times(2);
  
      control.replay();
  
      assertEquals(
          new JobDiff(ImmutableMap.of(), ImmutableSet.of(), ImmutableMap.of(0, task, 1, task)),
          JobDiff.compute(store, JOB, JobDiff.asMap(task, 2), NO_SCOPE));
      assertEquals(
          new JobDiff(ImmutableMap.of(), ImmutableSet.of(), ImmutableMap.of(0, task, 1, task)),
          JobDiff.compute(store, JOB, JobDiff.asMap(task, 2), CANARY_SCOPE));
    }
  
    @Test
    public void testInstancesAdded() {
      ITaskConfig task = makeTask("job", "echo");
  
      expectFetch(instance(task, 0), instance(task, 1)).times(2);
  
      control.replay();
  
      assertEquals(
          new JobDiff(ImmutableMap.of(), ImmutableSet.of(2, 3, 4), ImmutableMap.of(0, task, 1, task)),
          JobDiff.compute(store, JOB, JobDiff.asMap(task, 5), NO_SCOPE));
      assertEquals(
          new JobDiff(ImmutableMap.of(), ImmutableSet.of(), ImmutableMap.of(0, task, 1, task)),
          JobDiff.compute(store, JOB, JobDiff.asMap(task, 5), CANARY_SCOPE));
    }
  
    @Test
    public void testInstancesRemoved() {
      ITaskConfig task = makeTask("job", "echo");
  
      expectFetch(instance(task, 0), instance(task, 1), instance(task, 2)).times(2);
  
      control.replay();
  
      assertEquals(
          new JobDiff(ImmutableMap.of(1, task, 2, task), ImmutableSet.of(), ImmutableMap.of(0, task)),
          JobDiff.compute(store, JOB, JobDiff.asMap(task, 1), NO_SCOPE));
      assertEquals(
          new JobDiff(
              ImmutableMap.of(),
              ImmutableSet.of(),
              ImmutableMap.of(0, task, 1, task, 2, task)),
          JobDiff.compute(store, JOB, JobDiff.asMap(task, 1), CANARY_SCOPE));
    }
  
    @Test
    public void testFullUpdate() {
      ITaskConfig oldTask = makeTask("job", "echo");
      ITaskConfig newTask = makeTask("job", "echo2");
  
      expectFetch(instance(oldTask, 0), instance(oldTask, 1), instance(oldTask, 2)).times(2);
  
      control.replay();
  
      assertEquals(
          new JobDiff(
              ImmutableMap.of(0, oldTask, 1, oldTask, 2, oldTask),
              ImmutableSet.of(0, 1, 2),
              ImmutableMap.of()),
          JobDiff.compute(store, JOB, JobDiff.asMap(newTask, 3), NO_SCOPE));
      assertEquals(
          new JobDiff(
              ImmutableMap.of(0, oldTask),
              ImmutableSet.of(0),
              ImmutableMap.of(1, oldTask, 2, oldTask)),
          JobDiff.compute(store, JOB, JobDiff.asMap(newTask, 3), CANARY_SCOPE));
    }
  
    @Test
    public void testMultipleConfigsAndHoles() {
      ITaskConfig oldTask = makeTask("job", "echo");
      ITaskConfig oldTask2 = makeTask("job", "echo1");
      ITaskConfig newTask = makeTask("job", "echo2");
  
      expectFetch(
          instance(oldTask, 0), instance(newTask, 1), instance(oldTask2, 3), instance(oldTask, 4))
          .times(2);
  
      control.replay();
  
      assertEquals(
          new JobDiff(
              ImmutableMap.of(0, oldTask, 3, oldTask2, 4, oldTask),
              ImmutableSet.of(0, 2, 3, 4),
              ImmutableMap.of(1, newTask)),
          JobDiff.compute(store, JOB, JobDiff.asMap(newTask, 5), NO_SCOPE));
      assertEquals(
          new JobDiff(
              ImmutableMap.of(0, oldTask),
              ImmutableSet.of(0),
              ImmutableMap.of(1, newTask, 3, oldTask2, 4, oldTask)),
          JobDiff.compute(store, JOB, JobDiff.asMap(newTask, 5), CANARY_SCOPE));
    }
  
    @Test
    public void testUnchangedInstances() {
      ITaskConfig oldTask = makeTask("job", "echo");
      ITaskConfig newTask = makeTask("job", "echo2");
  
      expectFetch(instance(oldTask, 0), instance(newTask, 1), instance(oldTask, 2)).times(3);
  
      control.replay();
  
      assertEquals(
          new JobDiff(
              ImmutableMap.of(0, oldTask, 2, oldTask),
              ImmutableSet.of(0, 2),
              ImmutableMap.of(1, newTask)),
          JobDiff.compute(store, JOB, JobDiff.asMap(newTask, 3), NO_SCOPE));
      assertEquals(
          new JobDiff(
              ImmutableMap.of(0, oldTask),
              ImmutableSet.of(0),
              ImmutableMap.of(1, newTask, 2, oldTask)),
          JobDiff.compute(store, JOB, JobDiff.asMap(newTask, 3), CANARY_SCOPE));
      assertEquals(
          new JobDiff(
              ImmutableMap.of(0, oldTask),
              ImmutableSet.of(0),
              ImmutableMap.of(1, newTask, 2, oldTask)),
          JobDiff.compute(store, JOB, JobDiff.asMap(newTask, 4), CANARY_SCOPE));
    }
  
    @Test
    public void testObjectOverrides() {
      control.replay();
  
      JobDiff a = new JobDiff(
          ImmutableMap.of(0, makeTask("job", "echo")),
          ImmutableSet.of(0),
          ImmutableMap.of());
      JobDiff b = new JobDiff(
          ImmutableMap.of(0, makeTask("job", "echo")),
          ImmutableSet.of(0),
          ImmutableMap.of());
      JobDiff c = new JobDiff(
          ImmutableMap.of(0, makeTask("job", "echo1")),
          ImmutableSet.of(0),
          ImmutableMap.of());
      JobDiff d = new JobDiff(
          ImmutableMap.of(0, makeTask("job", "echo")),
          ImmutableSet.of(1),
          ImmutableMap.of());
      assertEquals(a, b);
      assertEquals(ImmutableSet.of(a), ImmutableSet.of(a, b));
      assertNotEquals(a, c);
      assertNotEquals(a, "a string");
      assertNotEquals(a, d);
      assertEquals(a.toString(), b.toString());
    }
  
    private IExpectationSetters<?> expectFetch(IAssignedTask... results) {
      ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder();
      for (IAssignedTask result : results) {
        tasks.add(IScheduledTask.build(new ScheduledTask().setAssignedTask(result.newBuilder())));
      }
  
      return expect(store.fetchTasks(Query.jobScoped(JOB).active()))
          .andReturn(tasks.build());
    }
  
    private static IAssignedTask instance(ITaskConfig config, int instance) {
      return IAssignedTask.build(
          new AssignedTask().setTask(config.newBuilder()).setInstanceId(instance));
    }
  
    private static ITaskConfig makeTask(String job, String config) {
      return ITaskConfig.build(new TaskConfig()
          .setJob(JobKeys.from("role", "test", job).newBuilder())
          .setOwner(new Identity().setUser("owner"))
          .setExecutorConfig(new ExecutorConfig().setData(config)));
    }
  }
  