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.http;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.stream.Collectors;
  
  import com.google.common.collect.ImmutableSet;
  
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.aurora.common.util.testing.FakeTicker;
  import org.apache.aurora.gen.JobKey;
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.scheduler.base.TaskGroupKey;
  import org.apache.aurora.scheduler.events.PubsubEvent;
  import org.apache.aurora.scheduler.filter.SchedulingFilter.Veto;
  import org.apache.aurora.scheduler.metadata.NearestFit;
  import org.apache.aurora.scheduler.scheduling.TaskGroup;
  import org.apache.aurora.scheduler.scheduling.TaskGroups;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.codehaus.jackson.JsonNode;
  import org.codehaus.jackson.map.ObjectMapper;
  import org.codehaus.jackson.node.ArrayNode;
  import org.codehaus.jackson.node.ObjectNode;
  import org.junit.Before;
  import org.junit.Test;
  
  import static org.easymock.EasyMock.expect;
  import static org.easymock.EasyMock.expectLastCall;
  import static org.easymock.EasyMock.replay;
  import static org.junit.Assert.assertEquals;
  
  public class PendingTasksTest extends EasyMockTest {
  
    private TaskGroups pendingTaskGroups;
    private NearestFit nearestFit;
  
    @Before
    public void setUp() {
      pendingTaskGroups = createMock(TaskGroups.class);
      nearestFit = new NearestFit(new FakeTicker());
    }
  
    /**
     * Create a {@link JsonNode} object to mimic the response.
     *
     * @param penaltyMs
     * @param taskIds
     * @param name
     * @param reasons
     * @return Json node for pending tasks whose values are initialized to the provided values.
     * @throws IOException
     */
    private JsonNode getMimicResponseJson(
        long penaltyMs, String[] taskIds, String name, List<String> reasons) throws IOException {
      ObjectMapper mapper = new ObjectMapper();
      ObjectNode mutablePendingTaskJson = mapper.createObjectNode();
      // Adding the key=value pairs to mutablePendingTaskJson.
      mutablePendingTaskJson.put("penaltyMs", penaltyMs);
      mutablePendingTaskJson.putArray("taskIds");
      for (String taskId : taskIds) {
        ((ArrayNode) mutablePendingTaskJson.get("taskIds")).add(taskId);
      }
      mutablePendingTaskJson.put("name", name);
      mutablePendingTaskJson.put("reason", reasons.toString());
      return mutablePendingTaskJson;
    }
  
    @Test
    public void testNoOffers() throws IOException {
      // Making a task that is not in PENDING state.
      IJobKey jobKey = IJobKey.build(new JobKey("role", "test", "nonPendingJob"));
      IScheduledTask task = TestUtils.makeTask(jobKey, "task0", 0,
          ScheduleStatus.ASSIGNED, 10, 10, 10);
  
      PubsubEvent.TaskStateChange taskStateChange = PubsubEvent.TaskStateChange.transition(
          task, ScheduleStatus.INIT);
  
      pendingTaskGroups.taskChangedState(taskStateChange);
      expectLastCall();
  
      // Recording the return value of pendingTaskGroups.getGroups().
      List<TaskGroup> taskGroupList = new ArrayList<>();
      expect(pendingTaskGroups.getGroups()).andReturn(taskGroupList).anyTimes();
  
      replay(pendingTaskGroups);
  
      // Testing.
      pendingTaskGroups.taskChangedState(taskStateChange);
      PendingTasks pendingTasks = new PendingTasks(pendingTaskGroups, nearestFit);
      JsonNode mimicResponseNoPendingTaskJson = new ObjectMapper().createArrayNode();
      JsonNode actualResponseJson = new ObjectMapper().valueToTree(
          pendingTasks.getOffers().getEntity());
      assertEquals(mimicResponseNoPendingTaskJson, actualResponseJson);
    }
  
    @Test
    public void testOffers() throws IOException {
      // Making pending tasks.
      IJobKey jobKey0 = IJobKey.build(new JobKey("role", "test", "jobA"));
      IJobKey jobKey1 = IJobKey.build(new JobKey("role", "test", "jobB"));
      // Task of jobA
      IScheduledTask task0 = TestUtils.makeTask(jobKey0, "task0", 0,
          ScheduleStatus.PENDING, 1000, 1000000, 10);
      // Tasks of jobB with two different TaskConfigs and thus two different TaskGroupKeys
      IScheduledTask task1 = TestUtils.makeTask(jobKey1, "task1", 0,
          ScheduleStatus.PENDING, 1000, 10, 1000000);
      IScheduledTask task2 = TestUtils.makeTask(jobKey1, "task2", 1,
          ScheduleStatus.PENDING, 100, 1, 100000);
  
      PubsubEvent.TaskStateChange taskStateChange0 = PubsubEvent.TaskStateChange.transition(
          task0, ScheduleStatus.INIT);
      PubsubEvent.TaskStateChange taskStateChange1 = PubsubEvent.TaskStateChange.transition(
          task1, ScheduleStatus.INIT);
      PubsubEvent.TaskStateChange taskStateChange2 = PubsubEvent.TaskStateChange.transition(
          task2, ScheduleStatus.INIT);
  
      pendingTaskGroups.taskChangedState(taskStateChange0);
      pendingTaskGroups.taskChangedState(taskStateChange1);
      pendingTaskGroups.taskChangedState(taskStateChange2);
      expectLastCall();
  
      // Recording the return value of pendingTaskGroups.getGroups().
      TaskGroupKey taskGroupKey0 = TaskGroupKey.from(task0.getAssignedTask().getTask());
      TaskGroupKey taskGroupKey1 = TaskGroupKey.from(task1.getAssignedTask().getTask());
      TaskGroupKey taskGroupKey2 = TaskGroupKey.from(task2.getAssignedTask().getTask());
      TaskGroup taskGroup0 = new TaskGroup(taskGroupKey0, "task0");
      TaskGroup taskGroup1 = new TaskGroup(taskGroupKey1, "task1");
      TaskGroup taskGroup2 = new TaskGroup(taskGroupKey2, "task2");
      List<TaskGroup> taskGroupList = new ArrayList<>();
      taskGroupList.add(taskGroup0);
      taskGroupList.add(taskGroup1);
      taskGroupList.add(taskGroup2);
      expect(pendingTaskGroups.getGroups()).andReturn(taskGroupList).anyTimes();
  
      // Creating vetoes for CPU and RAM, corresponding to task0.
      ImmutableSet<Veto> vetoes = ImmutableSet.<Veto>builder()
          .add(Veto.insufficientResources("CPU", 1))
          .add(Veto.insufficientResources("RAM", 1)).build();
      nearestFit.vetoed(taskGroupKey0, vetoes);
      // Creating vetoes for CPU and DISK, corresponding to task1.
      ImmutableSet<Veto> vetoes1 = ImmutableSet.<Veto>builder()
          .add(Veto.insufficientResources("CPU", 1))
          .add(Veto.insufficientResources("DISK", 1)).build();
      nearestFit.vetoed(taskGroupKey1, vetoes1);
        // Creating vetoes for CPU, corresponding to task2.
      ImmutableSet<Veto> vetoes2 = ImmutableSet.<Veto>builder()
          .add(Veto.insufficientResources("CPU", 1)).build();
      nearestFit.vetoed(taskGroupKey2, vetoes2);
      replay(pendingTaskGroups);
  
      // Testing.
      pendingTaskGroups.taskChangedState(taskStateChange0);
      pendingTaskGroups.taskChangedState(taskStateChange1);
      pendingTaskGroups.taskChangedState(taskStateChange2);
  
      PendingTasks pendingTasks = new PendingTasks(pendingTaskGroups, nearestFit);
      String[] taskIds0 = {"task0"};
      String[] taskIds1 = {"task1"};
      String[] taskIds2 = {"task2"};
      String[] reasonsArr0 = {"Insufficient: CPU", "Insufficient: RAM"};
      String[] reasonsArr1 = {"Insufficient: CPU", "Insufficient: DISK"};
      String[] reasonsArr2 = {"Insufficient: CPU"};
      List<String> reasons0 = Arrays.stream(reasonsArr0).collect(Collectors.toList());
      List<String> reasons1 = Arrays.stream(reasonsArr1).collect(Collectors.toList());
      List<String> reasons2 = Arrays.stream(reasonsArr2).collect(Collectors.toList());
  
      JsonNode mimicResponseTwoPendingTasksJson = new ObjectMapper().createArrayNode();
      JsonNode mimicJson0 = getMimicResponseJson(0, taskIds0, "role/test/jobA", reasons0);
      JsonNode mimicJson1 = getMimicResponseJson(0, taskIds1, "role/test/jobB", reasons1);
      JsonNode mimicJson2 = getMimicResponseJson(0, taskIds2, "role/test/jobB", reasons2);
      ((ArrayNode) mimicResponseTwoPendingTasksJson).add(mimicJson0);
      ((ArrayNode) mimicResponseTwoPendingTasksJson).add(mimicJson1);
      ((ArrayNode) mimicResponseTwoPendingTasksJson).add(mimicJson2);
  
      JsonNode actualResponseJson = new ObjectMapper().valueToTree(
          pendingTasks.getOffers().getEntity());
  
      assertEquals(mimicResponseTwoPendingTasksJson, actualResponseJson);
    }
  }
  