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 com.google.common.collect.ImmutableSet;
  
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.aurora.gen.InstanceKey;
  import org.apache.aurora.gen.InstanceTaskConfig;
  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.Range;
  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.base.TaskTestUtil;
  import org.apache.aurora.scheduler.state.StateManager;
  import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateInstructions;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
  import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
  import org.junit.Before;
  import org.junit.Test;
  
  public class AddTaskTest extends EasyMockTest {
    private static final IJobUpdateInstructions INSTRUCTIONS = IJobUpdateInstructions.build(
        new JobUpdateInstructions()
            .setDesiredState(new InstanceTaskConfig()
                .setTask(new TaskConfig())
                .setInstances(ImmutableSet.of(new Range(0, 0))))
            .setSettings(
                new JobUpdateSettings()
                    .setMinWaitInInstanceRunningMs(1000)));
    private static final IJobKey JOB = JobKeys.from("role", "env", "job");
    private static final IInstanceKey INSTANCE =
        IInstanceKey.build(new InstanceKey(JOB.newBuilder(), 0));
    private static final IJobUpdateKey UPDATE_ID =
            IJobUpdateKey.build(new JobUpdateKey(JOB.newBuilder(), "update_id"));
  
    private StorageTestUtil storageUtil;
    private StateManager stateManager;
    private InstanceActionHandler handler;
    private UpdateAgentReserver updateAgentReserver;
    private SlaKillController slaKillController;
  
    @Before
    public void setUp() {
      storageUtil = new StorageTestUtil(this);
      storageUtil.expectOperations();
      stateManager = createMock(StateManager.class);
      updateAgentReserver = createMock(UpdateAgentReserver.class);
      handler = new InstanceActionHandler.AddTask();
      slaKillController = createMock(SlaKillController.class);
    }
  
    @Test
    public void testAddInstance() throws Exception {
      storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active());
  
      stateManager.insertPendingTasks(
          storageUtil.mutableStoreProvider,
          INSTRUCTIONS.getDesiredState().getTask(),
          ImmutableSet.of(0));
  
      control.replay();
  
      handler.getReevaluationDelay(
          INSTANCE,
          INSTRUCTIONS,
          storageUtil.mutableStoreProvider,
          stateManager,
          updateAgentReserver,
          JobUpdateStatus.ROLLING_FORWARD,
          UPDATE_ID,
          slaKillController);
    }
  
    @Test
    public void testAddInstanceCollisionDoesNotThrow() throws Exception {
      storageUtil.expectTaskFetch(
          Query.instanceScoped(INSTANCE).active(),
          TaskTestUtil.makeTask("id", INSTANCE.getJobKey()));
  
      control.replay();
  
      handler.getReevaluationDelay(
          INSTANCE,
          INSTRUCTIONS,
          storageUtil.mutableStoreProvider,
          stateManager,
          updateAgentReserver,
          JobUpdateStatus.ROLLING_FORWARD,
          UPDATE_ID,
          slaKillController);
    }
  
    @Test(expected = IllegalStateException.class)
    public void testInstanceNotFound() throws Exception {
      storageUtil.expectTaskFetch(Query.instanceScoped(INSTANCE).active());
  
      control.replay();
  
      handler.getReevaluationDelay(
          INSTANCE,
          INSTRUCTIONS,
          storageUtil.mutableStoreProvider,
          stateManager,
          updateAgentReserver,
          JobUpdateStatus.ROLLING_BACK,
          UPDATE_ID,
          slaKillController);
    }
  }
  