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.base;
  
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSet;
  
  import org.apache.aurora.gen.AssignedTask;
  import org.apache.aurora.gen.Constraint;
  import org.apache.aurora.gen.Container;
  import org.apache.aurora.gen.Container._Fields;
  import org.apache.aurora.gen.DockerContainer;
  import org.apache.aurora.gen.DockerParameter;
  import org.apache.aurora.gen.Identity;
  import org.apache.aurora.gen.LimitConstraint;
  import org.apache.aurora.gen.MesosFetcherURI;
  import org.apache.aurora.gen.Metadata;
  import org.apache.aurora.gen.PartitionPolicy;
  import org.apache.aurora.gen.PercentageSlaPolicy;
  import org.apache.aurora.gen.Resource;
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.gen.ScheduledTask;
  import org.apache.aurora.gen.SlaPolicy;
  import org.apache.aurora.gen.TaskConfig;
  import org.apache.aurora.gen.TaskConstraint;
  import org.apache.aurora.gen.TaskEvent;
  import org.apache.aurora.gen.ValueConstraint;
  import org.apache.aurora.gen.apiConstants;
  import org.apache.aurora.scheduler.TierInfo;
  import org.apache.aurora.scheduler.TierManager;
  import org.apache.aurora.scheduler.TierManager.TierManagerImpl.TierConfig;
  import org.apache.aurora.scheduler.configuration.ConfigurationManager;
  import org.apache.aurora.scheduler.configuration.ConfigurationManager.ConfigurationManagerSettings;
  import org.apache.aurora.scheduler.configuration.executor.ExecutorConfig;
  import org.apache.aurora.scheduler.configuration.executor.ExecutorSettings;
  import org.apache.aurora.scheduler.filter.AttributeAggregate;
  import org.apache.aurora.scheduler.filter.SchedulingFilter.ResourceRequest;
  import org.apache.aurora.scheduler.storage.durability.ThriftBackfill;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
  import org.apache.mesos.v1.Protos;
  import org.apache.mesos.v1.Protos.ExecutorID;
  import org.apache.mesos.v1.Protos.ExecutorInfo;
  
  /**
   * Convenience methods for working with tasks.
   */
  public final class TaskTestUtil {
  
    public static final IJobKey JOB = JobKeys.from("role", "devel", "job");
    public static final TierInfo REVOCABLE_TIER =
        new TierInfo(true /* preemptible */, true /* revocable */);
    public static final TierInfo DEV_TIER =
        new TierInfo(true /* preemptible */, false /* revocable */);
    public static final TierInfo PREFERRED_TIER =
        new TierInfo(false /* preemptible */, false /* revocable */);
    public static final String REVOCABLE_TIER_NAME = "tier-revocable";
    public static final String PROD_TIER_NAME = "tier-prod";
    public static final String DEV_TIER_NAME = "tier-dev";
    public static final TierConfig TIER_CONFIG =
        new TierConfig(DEV_TIER_NAME, ImmutableMap.of(
            PROD_TIER_NAME, PREFERRED_TIER,
            DEV_TIER_NAME, DEV_TIER,
            REVOCABLE_TIER_NAME, REVOCABLE_TIER
        ));
    public static final TierManager TIER_MANAGER = new TierManager.TierManagerImpl(TIER_CONFIG);
    public static final ThriftBackfill THRIFT_BACKFILL = new ThriftBackfill();
    public static final ConfigurationManagerSettings CONFIGURATION_MANAGER_SETTINGS =
        new ConfigurationManagerSettings(
            ImmutableSet.of(_Fields.MESOS),
            false,
            ImmutableList.of(),
            true,
            true,
            true,
            true,
            2,
            1800L,
            ConfigurationManager.DEFAULT_ALLOWED_JOB_ENVIRONMENTS,
            false);
    public static final ExecutorID EXECUTOR_ID = ExecutorID.newBuilder()
        .setValue("PLACEHOLDER")
        .build();
    public static final ExecutorInfo EXECUTOR_INFO = ExecutorInfo.newBuilder()
        .setExecutorId(EXECUTOR_ID)
        .setName(apiConstants.AURORA_EXECUTOR_NAME)
        .setCommand(Protos.CommandInfo.newBuilder().build()).build();
    public static final ExecutorSettings EXECUTOR_SETTINGS = new ExecutorSettings(
        ImmutableMap.<String, ExecutorConfig>builder()
            .put(EXECUTOR_INFO.getName(), new ExecutorConfig(EXECUTOR_INFO, ImmutableList.of(), ""))
            .build(),
        false);
    public static final ConfigurationManager CONFIGURATION_MANAGER =
        new ConfigurationManager(CONFIGURATION_MANAGER_SETTINGS,
            TIER_MANAGER,
            THRIFT_BACKFILL,
            EXECUTOR_SETTINGS);
  
    private TaskTestUtil() {
      // Utility class.
    }
  
    public static ITaskConfig makeConfig(IJobKey job) {
      return makeConfig(job, true);
    }
  
    public static ITaskConfig makeConfig(IJobKey job, boolean prod) {
      return makeConfig(job, prod, Optional.of(SlaPolicy.percentageSlaPolicy(
          new PercentageSlaPolicy().setPercentage(95.0).setDurationSecs(1800))));
    }
  
    public static ITaskConfig makeConfig(IJobKey job, boolean prod, Optional<SlaPolicy> slaPolicy) {
      return ITaskConfig.build(new TaskConfig()
          .setJob(job.newBuilder())
          .setOwner(new Identity().setUser(job.getRole() + "-user"))
          .setIsService(true)
          .setPriority(1)
          .setMaxTaskFailures(-1)
          .setProduction(prod)
          .setTier(prod ? PROD_TIER_NAME : DEV_TIER_NAME)
          .setPartitionPolicy(new PartitionPolicy().setDelaySecs(5).setReschedule(true))
          .setSlaPolicy(slaPolicy.orElse(null))
          .setConstraints(ImmutableSet.of(
              new Constraint(
                  "valueConstraint",
                  TaskConstraint.value(
                      new ValueConstraint(true, ImmutableSet.of("value1", "value2")))),
              new Constraint(
                  "limitConstraint",
                  TaskConstraint.limit(new LimitConstraint(5)))))
          .setTaskLinks(ImmutableMap.of("http", "link", "admin", "otherLink"))
          .setContactEmail("foo@bar.com")
          .setMetadata(ImmutableSet.of(new Metadata("key", "value")))
          .setMesosFetcherUris(ImmutableSet.of(
              new MesosFetcherURI("pathA")
                  .setExtract(true)
                  .setCache(true)
                  .setOutputFile("test_1"),
              new MesosFetcherURI("pathB")
                  .setExtract(true)
                  .setCache(true)
                  .setOutputFile("test_2")))
          .setExecutorConfig(new org.apache.aurora.gen.ExecutorConfig(
              EXECUTOR_INFO.getName(),
              "config"))
          .setContainer(Container.docker(
              new DockerContainer("imagename")
                  .setParameters(ImmutableList.of(
                      new DockerParameter("a", "b"),
                      new DockerParameter("c", "d")))))
          .setResources(ImmutableSet.of(
              Resource.numCpus(1.0),
              Resource.ramMb(1024),
              Resource.diskMb(1024),
              Resource.namedPort("http"))));
    }
  
    public static IScheduledTask makeTask(String id, IJobKey job) {
      return makeTask(id, makeConfig(job));
    }
  
    public static IScheduledTask makeTask(String id, IJobKey job, int instanceId) {
      return makeTask(id, makeConfig(job), instanceId, Optional.empty());
    }
  
    public static IScheduledTask makeTask(String id, IJobKey job, int instanceId, boolean prod) {
      return makeTask(id, makeConfig(job, prod), instanceId, Optional.empty());
    }
  
    public static IScheduledTask makeTask(String id, IJobKey job, int instanceId, String agentId) {
      return makeTask(id, makeConfig(job), instanceId, Optional.of(agentId));
    }
  
    public static IScheduledTask makeTask(String id, ITaskConfig config) {
      return makeTask(id, config, 2);
    }
  
    public static IScheduledTask makeTask(String id, ITaskConfig config, int instanceId) {
      return makeTask(id, config, instanceId, Optional.empty());
    }
  
    public static IScheduledTask makeTask(
        String id,
        ITaskConfig config,
        int instanceId,
        Optional<String> agentId) {
  
      AssignedTask assignedTask = new AssignedTask()
          .setInstanceId(instanceId)
          .setTaskId(id)
          .setAssignedPorts(ImmutableMap.of("http", 1000))
          .setTask(config.newBuilder());
      if (agentId.isPresent()) {
        assignedTask.setSlaveId(agentId.get());
      }
  
      return IScheduledTask.build(new ScheduledTask()
          .setStatus(ScheduleStatus.ASSIGNED)
          .setTaskEvents(ImmutableList.of(
              new TaskEvent(100L, ScheduleStatus.PENDING)
                  .setMessage("message")
                  .setScheduler("scheduler"),
              new TaskEvent(101L, ScheduleStatus.ASSIGNED)
                  .setMessage("message")
                  .setScheduler("scheduler2")))
          .setAncestorId("ancestor")
          .setFailureCount(3)
          .setTimesPartitioned(2)
          .setAssignedTask(assignedTask));
    }
  
    public static IScheduledTask addStateTransition(
        IScheduledTask task,
        ScheduleStatus status,
        long timestamp) {
  
      ScheduledTask builder = task.newBuilder();
      builder.setStatus(status);
      builder.addToTaskEvents(new TaskEvent()
          .setTimestamp(timestamp)
          .setStatus(status)
          .setScheduler("scheduler"));
      return IScheduledTask.build(builder);
    }
  
    public static String tierConfigFile() {
      return "{\"default\": \"preemptible\","
          + "\"tiers\":{"
          + "\"preferred\": {\"revocable\": false, \"preemptible\": false},"
          + "\"preemptible\": {\"revocable\": false, \"preemptible\": true},"
          + "\"revocable\": {\"revocable\": true, \"preemptible\": true}"
          + "}}";
    }
  
    public static Map<String, TierInfo> tierInfos() {
      return ImmutableMap.of(
          "preferred", PREFERRED_TIER,
          "preemptible", DEV_TIER,
          "revocable", REVOCABLE_TIER);
    }
  
    public static Set<org.apache.aurora.gen.TierConfig> tierConfigs() {
      return ImmutableSet.of(
          new org.apache.aurora.gen.TierConfig("preferred", PREFERRED_TIER.toMap()),
          new org.apache.aurora.gen.TierConfig("preemptible", DEV_TIER.toMap()),
          new org.apache.aurora.gen.TierConfig("revocable", REVOCABLE_TIER.toMap())
      );
    }
  
    public static ResourceRequest toResourceRequest(ITaskConfig task) {
      return ResourceRequest.fromTask(
          task,
          EXECUTOR_SETTINGS,
          AttributeAggregate.empty(),
          TIER_MANAGER);
    }
  }
  