aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ma...@apache.org
Subject aurora git commit: AURORA-1624 New thrift API for retrieving tier configuration
Date Tue, 07 Jun 2016 20:45:00 GMT
Repository: aurora
Updated Branches:
  refs/heads/master e6674939c -> 8cce5bc00


AURORA-1624 New thrift API for retrieving tier configuration

Bugs closed: AURORA-1624

Reviewed at https://reviews.apache.org/r/48082/


Project: http://git-wip-us.apache.org/repos/asf/aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/8cce5bc0
Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/8cce5bc0
Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/8cce5bc0

Branch: refs/heads/master
Commit: 8cce5bc00d6b14113c428ae885e8facfe58bd985
Parents: e667493
Author: Mehrdad Nurolahzade <mehrdad@nurolahzade.com>
Authored: Tue Jun 7 13:41:48 2016 -0700
Committer: Maxim Khutornenko <maxim@apache.org>
Committed: Tue Jun 7 13:41:48 2016 -0700

----------------------------------------------------------------------
 RELEASE-NOTES.md                                |  6 ++
 .../thrift/org/apache/aurora/gen/api.thrift     | 21 ++++++
 .../org/apache/aurora/scheduler/TierInfo.java   | 14 ++++
 .../apache/aurora/scheduler/TierManager.java    | 41 +++++++++-
 .../aurora/scheduler/base/TaskTestUtil.java     | 21 +++++-
 .../scheduler/thrift/ReadOnlySchedulerImpl.java | 21 +++++-
 .../thrift/SchedulerThriftInterface.java        |  5 ++
 .../org/apache/aurora/scheduler/tiers.json      |  1 +
 .../aurora/scheduler/TierManagerTest.java       | 79 +++++++++++++++++++-
 .../apache/aurora/scheduler/TierModuleTest.java |  6 +-
 .../thrift/ReadOnlySchedulerImplTest.java       | 43 ++++++++++-
 11 files changed, 248 insertions(+), 10 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/RELEASE-NOTES.md
----------------------------------------------------------------------
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 96bc3ca..741c725 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -42,6 +42,10 @@
 - Experimental support for a webhook feature which POSTs all task state changes to a user
defined
   endpoint.
 
+- Added support for specifying the default tier name in tier configuration file (`tiers.json`).
The
+  `default` property is required and is initialized with the `preemptible` tier (`preemptible`
tier
+  tasks can be preempted but their resources cannot be revoked).
+
 ### Deprecations and removals:
 
 - Deprecated `--restart-threshold` option in the `aurora job restart` command to match the
job
@@ -57,6 +61,8 @@
 - Deprecated `numCpus`, `ramMb` and `diskMb` fields in `TaskConfig` and `ResourceAggregate`
thrift
   structs. Use `set<Resource> resources` to specify task resources or quota values.
 - The endpoint `/slaves` is deprecated. Please use `/agents` instead.
+- Deprecated `production` field in `TaskConfig` thrift struct. Use `tier` field to specify
task
+  scheduling and resource handling behavior.
 
 0.13.0
 ------

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/api/src/main/thrift/org/apache/aurora/gen/api.thrift
----------------------------------------------------------------------
diff --git a/api/src/main/thrift/org/apache/aurora/gen/api.thrift b/api/src/main/thrift/org/apache/aurora/gen/api.thrift
index ed94f24..3e6daf4 100644
--- a/api/src/main/thrift/org/apache/aurora/gen/api.thrift
+++ b/api/src/main/thrift/org/apache/aurora/gen/api.thrift
@@ -242,6 +242,7 @@ struct TaskConfig {
  10: i64 diskMb
  11: i32 priority
  13: i32 maxTaskFailures
+ // TODO(mnurolahzade): Deprecated. See AURORA-1708.
  /** Whether this is a production task, which can preempt. */
  18: optional bool production
  /** Task tier type. */
@@ -905,6 +906,22 @@ struct GetJobUpdateDiffResult {
   4: set<ConfigGroup> unchanged
 }
 
+/** Tier information. */
+struct TierConfig {
+  /** Name of tier. */
+  1: string name
+  /** Tier attributes. */
+  2: map<string, string> settings
+}
+
+/** Result of the getTierConfigResult call. */
+struct GetTierConfigResult {
+  /** Name of the default tier. */
+  1: string defaultTierName
+  /** Set of tier configurations. */
+  2: set<TierConfig> tiers
+}
+
 /** Information about the scheduler. */
 struct ServerInfo {
   1: string clusterName
@@ -932,6 +949,7 @@ union Result {
   24: GetJobUpdateDetailsResult getJobUpdateDetailsResult
   25: PulseJobUpdateResult pulseJobUpdateResult
   26: GetJobUpdateDiffResult getJobUpdateDiffResult
+  27: GetTierConfigResult getTierConfigResult
 }
 
 struct ResponseDetail {
@@ -996,6 +1014,9 @@ service ReadOnlyScheduler {
 
   /** Gets the diff between client (desired) and server (current) job states. */
   Response getJobUpdateDiff(1: JobUpdateRequest request)
+
+  /** Gets tier configurations. */
+  Response getTierConfigs()
 }
 
 service AuroraSchedulerManager extends ReadOnlyScheduler {

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/main/java/org/apache/aurora/scheduler/TierInfo.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/TierInfo.java b/src/main/java/org/apache/aurora/scheduler/TierInfo.java
index ac8901f..c45b949 100644
--- a/src/main/java/org/apache/aurora/scheduler/TierInfo.java
+++ b/src/main/java/org/apache/aurora/scheduler/TierInfo.java
@@ -13,9 +13,11 @@
  */
 package org.apache.aurora.scheduler;
 
+import java.util.Map;
 import java.util.Objects;
 
 import com.google.common.base.MoreObjects;
+import com.google.common.collect.ImmutableMap;
 
 import org.codehaus.jackson.annotate.JsonCreator;
 import org.codehaus.jackson.annotate.JsonProperty;
@@ -56,6 +58,18 @@ public final class TierInfo {
     return revocable;
   }
 
+  /**
+   * Gets the map of tier attribute names to values.
+   *
+   * @return A readonly view of all tier attributes.
+   */
+  public Map<String, String> toMap() {
+    return ImmutableMap.of(
+        "preemptible", String.valueOf(preemptible),
+        "revocable", String.valueOf(revocable)
+    );
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(preemptible, revocable);

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/main/java/org/apache/aurora/scheduler/TierManager.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/TierManager.java b/src/main/java/org/apache/aurora/scheduler/TierManager.java
index b96189b..af54cab 100644
--- a/src/main/java/org/apache/aurora/scheduler/TierManager.java
+++ b/src/main/java/org/apache/aurora/scheduler/TierManager.java
@@ -18,6 +18,7 @@ import java.util.Map;
 import javax.inject.Inject;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMap;
 
 import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
@@ -42,21 +43,49 @@ public interface TierManager {
    */
   TierInfo getTier(ITaskConfig taskConfig);
 
+  /**
+   * Gets name of the default tier.
+   *
+   * @return Name for the default tier.
+   */
+  String getDefaultTierName();
+
+  /**
+   * Gets the map of tier name to {@link TierInfo} instance.
+   *
+   * @return A readonly view of all tiers.
+   */
+  Map<String, TierInfo> getTiers();
+
   class TierManagerImpl implements TierManager {
     private final TierConfig tierConfig;
 
     @VisibleForTesting
     public static class TierConfig {
+      private final String defaultTier;
       private final Map<String, TierInfo> tiers;
 
       @VisibleForTesting
       @JsonCreator
-      public TierConfig(@JsonProperty("tiers") Map<String, TierInfo> tiers) {
+      public TierConfig(
+          @JsonProperty("default") String defaultTier,
+          @JsonProperty("tiers") Map<String, TierInfo> tiers) {
+
+        checkArgument(!Strings.isNullOrEmpty(defaultTier), "Default tier name cannot be empty.");
         checkArgument(!tiers.isEmpty(), "Tiers cannot be empty.");
+        checkArgument(
+            tiers.containsKey(defaultTier),
+            "Default tier name should match supplied tiers.");
+        this.defaultTier = defaultTier;
         this.tiers = ImmutableMap.copyOf(tiers);
       }
 
       @VisibleForTesting
+      public String getDefault() {
+        return defaultTier;
+      }
+
+      @VisibleForTesting
       public Map<String, TierInfo> getTiers() {
         return tiers;
       }
@@ -82,5 +111,15 @@ public interface TierManager {
               .orElseThrow(() -> new IllegalStateException(
                   format("No matching implicit tier for task of job %s", taskConfig.getJob())));
     }
+
+    @Override
+    public String getDefaultTierName() {
+      return tierConfig.defaultTier;
+    }
+
+    @Override
+    public Map<String, TierInfo> getTiers() {
+      return tierConfig.tiers;
+    }
   }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java b/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
index e431b58..3ea0992 100644
--- a/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
+++ b/src/main/java/org/apache/aurora/scheduler/base/TaskTestUtil.java
@@ -13,6 +13,8 @@
  */
 package org.apache.aurora.scheduler.base;
 
+import java.util.Map;
+
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
@@ -61,8 +63,23 @@ public final class TaskTestUtil {
       new TierInfo(false /* preemptible */, false /* revocable */);
   public static final String DEV_TIER_NAME = "tier-dev";
   public static final TierConfig DEV_TIER_CONFIG =
-      new TierConfig(ImmutableMap.of(DEV_TIER_NAME, DEV_TIER));
-  public static final TierManager DEV_TIER_MANAGER = taskConfig -> DEV_TIER;
+      new TierConfig(DEV_TIER_NAME, ImmutableMap.of(DEV_TIER_NAME, DEV_TIER));
+  public static final TierManager DEV_TIER_MANAGER = new TierManager() {
+    @Override
+    public TierInfo getTier(ITaskConfig taskConfig) {
+      return DEV_TIER;
+    }
+
+    @Override
+    public String getDefaultTierName() {
+      return DEV_TIER_NAME;
+    }
+
+    @Override
+    public Map<String, TierInfo> getTiers() {
+      return ImmutableMap.of(DEV_TIER_NAME, DEV_TIER);
+    }
+  };
   public static final ConfigurationManagerSettings CONFIGURATION_MANAGER_SETTINGS =
       new ConfigurationManagerSettings(
           ImmutableSet.of(_Fields.MESOS),

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java b/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
index 0d4f044..4202f0e 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImpl.java
@@ -40,6 +40,7 @@ import com.google.common.collect.Multimap;
 import com.google.common.collect.Multimaps;
 import com.google.common.collect.Sets;
 
+import org.apache.aurora.GuavaUtils;
 import org.apache.aurora.common.base.MorePreconditions;
 import org.apache.aurora.gen.ConfigGroup;
 import org.apache.aurora.gen.ConfigSummary;
@@ -50,6 +51,7 @@ import org.apache.aurora.gen.GetJobUpdateSummariesResult;
 import org.apache.aurora.gen.GetJobsResult;
 import org.apache.aurora.gen.GetPendingReasonResult;
 import org.apache.aurora.gen.GetQuotaResult;
+import org.apache.aurora.gen.GetTierConfigResult;
 import org.apache.aurora.gen.JobConfiguration;
 import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.JobSummary;
@@ -69,6 +71,8 @@ import org.apache.aurora.gen.ScheduleStatusResult;
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.gen.TaskQuery;
+import org.apache.aurora.gen.TierConfig;
+import org.apache.aurora.scheduler.TierManager;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Jobs;
 import org.apache.aurora.scheduler.base.Query;
@@ -119,6 +123,7 @@ class ReadOnlySchedulerImpl implements ReadOnlyScheduler.Iface {
   private final NearestFit nearestFit;
   private final CronPredictor cronPredictor;
   private final QuotaManager quotaManager;
+  private final TierManager tierManager;
 
   @Inject
   ReadOnlySchedulerImpl(
@@ -126,13 +131,15 @@ class ReadOnlySchedulerImpl implements ReadOnlyScheduler.Iface {
       Storage storage,
       NearestFit nearestFit,
       CronPredictor cronPredictor,
-      QuotaManager quotaManager) {
+      QuotaManager quotaManager,
+      TierManager tierManager) {
 
     this.configurationManager = requireNonNull(configurationManager);
     this.storage = requireNonNull(storage);
     this.nearestFit = requireNonNull(nearestFit);
     this.cronPredictor = requireNonNull(cronPredictor);
     this.quotaManager = requireNonNull(quotaManager);
+    this.tierManager = requireNonNull(tierManager);
   }
 
   @Override
@@ -368,6 +375,18 @@ class ReadOnlySchedulerImpl implements ReadOnlyScheduler.Iface {
     });
   }
 
+  @Override
+  public Response getTierConfigs() throws TException {
+    return ok(Result.getTierConfigResult(
+        new GetTierConfigResult(
+            tierManager.getDefaultTierName(),
+            tierManager.getTiers()
+                .entrySet()
+                .stream()
+                .map(entry -> new TierConfig(entry.getKey(), entry.getValue().toMap()))
+                .collect(GuavaUtils.toImmutableSet()))));
+  }
+
   private static Set<ConfigGroup> instancesToConfigGroups(Map<Integer, ITaskConfig>
tasks) {
     Multimap<ITaskConfig, Integer> instancesByDetails = Multimaps.invertFrom(
         Multimaps.forMap(tasks),

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
index 4d032b9..b534abf 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
@@ -396,6 +396,11 @@ class SchedulerThriftInterface implements AnnotatedAuroraAdmin {
     return readOnlyScheduler.getJobUpdateDiff(request);
   }
 
+  @Override
+  public Response getTierConfigs() throws TException {
+    return readOnlyScheduler.getTierConfigs();
+  }
+
   private void validateLockForTasks(Iterable<IScheduledTask> tasks) throws LockException
{
     ImmutableSet<IJobKey> uniqueKeys = FluentIterable.from(tasks)
         .transform(Tasks::getJob)

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/main/resources/org/apache/aurora/scheduler/tiers.json
----------------------------------------------------------------------
diff --git a/src/main/resources/org/apache/aurora/scheduler/tiers.json b/src/main/resources/org/apache/aurora/scheduler/tiers.json
index f724c5a..34ddb1d 100644
--- a/src/main/resources/org/apache/aurora/scheduler/tiers.json
+++ b/src/main/resources/org/apache/aurora/scheduler/tiers.json
@@ -1,4 +1,5 @@
 {
+  "default": "preemptible",
   "tiers":
   {
     "preferred":

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/test/java/org/apache/aurora/scheduler/TierManagerTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/TierManagerTest.java b/src/test/java/org/apache/aurora/scheduler/TierManagerTest.java
index 95174bb..d4b71f8 100644
--- a/src/test/java/org/apache/aurora/scheduler/TierManagerTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/TierManagerTest.java
@@ -13,19 +13,29 @@
  */
 package org.apache.aurora.scheduler;
 
+import java.util.Map;
+
+import com.google.common.collect.ImmutableMap;
+
 import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.scheduler.TierManager.TierManagerImpl;
-import org.apache.aurora.scheduler.base.TaskTestUtil;
 import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
 import org.junit.Test;
 
 import static org.apache.aurora.scheduler.TierModule.parseTierConfig;
+import static org.apache.aurora.scheduler.base.TaskTestUtil.DEV_TIER;
+import static org.apache.aurora.scheduler.base.TaskTestUtil.PREFERRED_TIER;
 import static org.apache.aurora.scheduler.base.TaskTestUtil.REVOCABLE_TIER;
 import static org.junit.Assert.assertEquals;
 
 public class TierManagerTest {
+  private static final Map<String, TierInfo> TIERS = ImmutableMap.of(
+      "preferred", PREFERRED_TIER,
+      "preemptible", DEV_TIER,
+      "revocable", REVOCABLE_TIER);
   private static final TierManager TIER_MANAGER = new TierManagerImpl(
-      parseTierConfig("{\"tiers\":{"
+      parseTierConfig("{\"default\": \"preemptible\","
+          + "\"tiers\":{"
           + "\"preferred\": {\"revocable\": false, \"preemptible\": false},"
           + "\"preemptible\": {\"revocable\": false, \"preemptible\": true},"
           + "\"revocable\": {\"revocable\": true, \"preemptible\": true}"
@@ -38,13 +48,76 @@ public class TierManagerTest {
         TIER_MANAGER.getTier(ITaskConfig.build(new TaskConfig().setTier("revocable"))));
   }
 
+  @Test
+  public void testRevocableAndProduction() {
+    assertEquals(
+        REVOCABLE_TIER,
+        TIER_MANAGER.getTier(ITaskConfig.build(new TaskConfig()
+            .setTier("revocable")
+            .setProduction(true))));
+  }
+
+  @Test
+  public void testPreemptibleAndProduction() {
+    assertEquals(
+        DEV_TIER,
+        TIER_MANAGER.getTier(ITaskConfig.build(new TaskConfig()
+            .setTier("preemptible")
+            .setProduction(true))));
+  }
+
   @Test(expected = IllegalArgumentException.class)
   public void testNameMismatch() {
     TIER_MANAGER.getTier(ITaskConfig.build(new TaskConfig().setTier("Revocable")));
   }
 
   @Test
+  public void testProduction() {
+    assertEquals(
+        PREFERRED_TIER,
+        TIER_MANAGER.getTier(ITaskConfig.build(new TaskConfig().setProduction(true))));
+  }
+
+  @Test
   public void testNoTierInTaskConfig() {
-    assertEquals(TaskTestUtil.DEV_TIER, TIER_MANAGER.getTier(ITaskConfig.build(new TaskConfig())));
+    assertEquals(DEV_TIER, TIER_MANAGER.getTier(ITaskConfig.build(new TaskConfig())));
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void testBadTierConfiguration() {
+    TierManager tierManager = new TierManagerImpl(
+        parseTierConfig("{\"default\": \"revocable\","
+            + "\"tiers\":{"
+            + "\"preferred\": {\"revocable\": false, \"preemptible\": false},"
+            + "\"revocable\": {\"revocable\": true, \"preemptible\": true}"
+            + "}}"));
+    // preemptible: false, revocable: false
+    ITaskConfig taskConfig = ITaskConfig.build(new TaskConfig());
+    tierManager.getTier(taskConfig);
+  }
+
+  @Test
+  public void testDefault() {
+    assertEquals("preemptible", TIER_MANAGER.getDefaultTierName());
+  }
+
+  @Test
+  public void testTiers() {
+    assertEquals(TIERS, TIER_MANAGER.getTiers());
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testMissingDefaultInTierConfig() {
+    new TierManagerImpl.TierConfig(null, TIERS);
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testEmptyTiersInTierConfig() {
+    new TierManagerImpl.TierConfig("preemptible", ImmutableMap.of());
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testMismatchingDefaultInTierConfig() {
+    new TierManagerImpl.TierConfig("something", TIERS);
   }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/test/java/org/apache/aurora/scheduler/TierModuleTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/TierModuleTest.java b/src/test/java/org/apache/aurora/scheduler/TierModuleTest.java
index 6b4e7a0..58d95dc 100644
--- a/src/test/java/org/apache/aurora/scheduler/TierModuleTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/TierModuleTest.java
@@ -52,7 +52,8 @@ public class TierModuleTest {
     assertEquals(
         ImmutableMap.of("revocable", REVOCABLE_TIER),
         parseTierConfig(
-            "{\"tiers\":{"
+            "{\"default\": \"revocable\","
+            + "\"tiers\":{"
             + "\"revocable\": {\"revocable\": true, \"preemptible\": true}"
             + "}}")
             .getTiers());
@@ -60,6 +61,7 @@ public class TierModuleTest {
 
   @Test(expected = RuntimeException.class)
   public void testTierConfigExtraKeysNotAllowed() {
-    parseTierConfig("{\"tiers\":{\"revocable\": {\"revocable\": true, \"foo\": false}}}");
+    parseTierConfig("{\"default\": \"revocable\", "
+        + "\"tiers\":{\"revocable\": {\"revocable\": true, \"foo\": false}}}");
   }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/8cce5bc0/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
b/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
index 2122f74..6f355d6 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/ReadOnlySchedulerImplTest.java
@@ -15,6 +15,7 @@ package org.apache.aurora.scheduler.thrift;
 
 import java.util.Date;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import javax.annotation.Nullable;
@@ -23,6 +24,7 @@ import com.google.common.base.Function;
 import com.google.common.base.Optional;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
@@ -35,6 +37,7 @@ import org.apache.aurora.gen.ConfigSummary;
 import org.apache.aurora.gen.ConfigSummaryResult;
 import org.apache.aurora.gen.GetJobUpdateDiffResult;
 import org.apache.aurora.gen.GetQuotaResult;
+import org.apache.aurora.gen.GetTierConfigResult;
 import org.apache.aurora.gen.Identity;
 import org.apache.aurora.gen.JobConfiguration;
 import org.apache.aurora.gen.JobKey;
@@ -61,6 +64,9 @@ import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.TaskConfig;
 import org.apache.aurora.gen.TaskQuery;
+import org.apache.aurora.gen.TierConfig;
+import org.apache.aurora.scheduler.TierInfo;
+import org.apache.aurora.scheduler.TierManager;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.Query.Builder;
@@ -91,6 +97,9 @@ import org.junit.Test;
 import static org.apache.aurora.gen.ResponseCode.INVALID_REQUEST;
 import static org.apache.aurora.scheduler.base.Numbers.convertRanges;
 import static org.apache.aurora.scheduler.base.Numbers.toRanges;
+import static org.apache.aurora.scheduler.base.TaskTestUtil.DEV_TIER;
+import static org.apache.aurora.scheduler.base.TaskTestUtil.PREFERRED_TIER;
+import static org.apache.aurora.scheduler.base.TaskTestUtil.REVOCABLE_TIER;
 import static org.apache.aurora.scheduler.resources.ResourceBag.LARGE;
 import static org.apache.aurora.scheduler.resources.ResourceBag.MEDIUM;
 import static org.apache.aurora.scheduler.resources.ResourceBag.SMALL;
@@ -122,6 +131,7 @@ public class ReadOnlySchedulerImplTest extends EasyMockTest {
   private NearestFit nearestFit;
   private CronPredictor cronPredictor;
   private QuotaManager quotaManager;
+  private TierManager tierManager;
 
   private ReadOnlyScheduler.Iface thrift;
 
@@ -132,13 +142,15 @@ public class ReadOnlySchedulerImplTest extends EasyMockTest {
     nearestFit = createMock(NearestFit.class);
     cronPredictor = createMock(CronPredictor.class);
     quotaManager = createMock(QuotaManager.class);
+    tierManager = createMock(TierManager.class);
 
     thrift = new ReadOnlySchedulerImpl(
         TaskTestUtil.CONFIGURATION_MANAGER,
         storageUtil.storage,
         nearestFit,
         cronPredictor,
-        quotaManager);
+        quotaManager,
+        tierManager);
   }
 
   @Test
@@ -844,4 +856,33 @@ public class ReadOnlySchedulerImplTest extends EasyMockTest {
         .setConfig(task)
         .setInstances(ImmutableSet.of(range));
   }
+
+  @Test
+  public void testGetTierConfig() throws Exception {
+    expect(tierManager.getDefaultTierName()).andReturn("preemptible");
+    expect(tierManager.getTiers()).andReturn(tierInfos());
+    control.replay();
+
+    GetTierConfigResult expected = new GetTierConfigResult()
+        .setDefaultTierName("preemptible")
+        .setTiers(tierConfigs());
+
+    Response response = assertOkResponse(thrift.getTierConfigs());
+    assertEquals(expected, response.getResult().getGetTierConfigResult());
+  }
+
+  private static Map<String, TierInfo> tierInfos() {
+    return ImmutableMap.of(
+        "preferred", PREFERRED_TIER,
+        "preemptible", DEV_TIER,
+        "revocable", REVOCABLE_TIER);
+  }
+
+  private static Set<TierConfig> tierConfigs() {
+    return ImmutableSet.of(
+        new TierConfig("preferred", PREFERRED_TIER.toMap()),
+        new TierConfig("preemptible", DEV_TIER.toMap()),
+        new TierConfig("revocable", REVOCABLE_TIER.toMap())
+    );
+  }
 }


Mime
View raw message