Allow allocation token discounts on premiums and for multiple years (#744)

* Allow allocation token discounts on premiums and for multiple years

* Add domain check flow tests

* Address code review comments

* Update schema file

core/src/main/java/google/registry/flows/domain/DomainPricingLogic.java

@@ -60,7 +60,7 @@ public final class DomainPricingLogic {
    * <p>If {@code allocationToken} is present and the domain is non-premium, that discount will be
    * applied to the first year.
    */
-  public FeesAndCredits getCreatePrice(
+  FeesAndCredits getCreatePrice(
       Registry registry,
       String domainName,
       DateTime dateTime,
@@ -104,8 +104,8 @@ public final class DomainPricingLogic {
 
   /** Returns a new renew price for the pricer. */
   @SuppressWarnings("unused")
-  public FeesAndCredits getRenewPrice(
+  FeesAndCredits getRenewPrice(Registry registry, String domainName, DateTime dateTime, int years)
-      Registry registry, String domainName, DateTime dateTime, int years) throws EppException {
+      throws EppException {
     DomainPrices domainPrices = getPricesForDomainName(domainName, dateTime);
     BigDecimal renewCost = domainPrices.getRenewCost().multipliedBy(years).getAmount();
     return customLogic.customizeRenewPrice(
@@ -123,7 +123,7 @@ public final class DomainPricingLogic {
   }
 
   /** Returns a new restore price for the pricer. */
-  public FeesAndCredits getRestorePrice(
+  FeesAndCredits getRestorePrice(
       Registry registry, String domainName, DateTime dateTime, boolean isExpired)
       throws EppException {
     DomainPrices domainPrices = getPricesForDomainName(domainName, dateTime);
@@ -147,7 +147,7 @@ public final class DomainPricingLogic {
   }
 
   /** Returns a new transfer price for the pricer. */
-  public FeesAndCredits getTransferPrice(Registry registry, String domainName, DateTime dateTime)
+  FeesAndCredits getTransferPrice(Registry registry, String domainName, DateTime dateTime)
       throws EppException {
     DomainPrices domainPrices = getPricesForDomainName(domainName, dateTime);
     return customLogic.customizeTransferPrice(
@@ -168,7 +168,7 @@ public final class DomainPricingLogic {
   }
 
   /** Returns a new update price for the pricer. */
-  public FeesAndCredits getUpdatePrice(Registry registry, String domainName, DateTime dateTime)
+  FeesAndCredits getUpdatePrice(Registry registry, String domainName, DateTime dateTime)
       throws EppException {
     CurrencyUnit currency = registry.getCurrency();
     BaseFee feeOrCredit = Fee.create(zeroInCurrency(currency), FeeType.UPDATE, false);
@@ -191,16 +191,20 @@ public final class DomainPricingLogic {
       throws EppException {
     if (allocationToken.isPresent()
         && allocationToken.get().getDiscountFraction() != 0.0
-        && domainPrices.isPremium()) {
+        && domainPrices.isPremium()
+        && !allocationToken.get().shouldDiscountPremiums()) {
       throw new AllocationTokenInvalidForPremiumNameException();
     }
     Money oneYearCreateCost = domainPrices.getCreateCost();
     Money totalDomainCreateCost = oneYearCreateCost.multipliedBy(years);
-    // If a discount is applicable, apply it only to the first year
+
+    // Apply the allocation token discount, if applicable.
     if (allocationToken.isPresent()) {
+      int discountedYears = Math.min(years, allocationToken.get().getDiscountYears());
       Money discount =
           oneYearCreateCost.multipliedBy(
-              allocationToken.get().getDiscountFraction(), RoundingMode.HALF_UP);
+              discountedYears * allocationToken.get().getDiscountFraction(),
+              RoundingMode.HALF_EVEN);
       totalDomainCreateCost = totalDomainCreateCost.minus(discount);
     }
     return totalDomainCreateCost;
@@ -209,7 +213,7 @@ public final class DomainPricingLogic {
   /** An allocation token was provided that is invalid for premium domains. */
   public static class AllocationTokenInvalidForPremiumNameException
       extends CommandUseErrorException {
-    public AllocationTokenInvalidForPremiumNameException() {
+    AllocationTokenInvalidForPremiumNameException() {
       super("A nonzero discount code cannot be applied to premium domains");
     }
   }

core/src/main/java/google/registry/model/BackupGroupRoot.java

@@ -14,7 +14,6 @@
 
 package google.registry.model;
 
-import com.google.common.annotations.VisibleForTesting;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
 import javax.persistence.MappedSuperclass;
@@ -30,6 +29,7 @@ import javax.xml.bind.annotation.XmlTransient;
  */
 @MappedSuperclass
 public abstract class BackupGroupRoot extends ImmutableObject {
+
   /**
    * An automatically managed timestamp of when this object was last written to Datastore.
    *
@@ -40,7 +40,6 @@ public abstract class BackupGroupRoot extends ImmutableObject {
   // Prevents subclasses from unexpectedly accessing as property (e.g., HostResource), which would
   // require an unnecessary non-private setter method.
   @Access(AccessType.FIELD)
-  @VisibleForTesting
   UpdateAutoTimestamp updateTimestamp = UpdateAutoTimestamp.create(null);
 
   /** Get the {@link UpdateAutoTimestamp} for this entity. */

core/src/main/java/google/registry/model/domain/DomainContent.java

@@ -534,8 +534,10 @@ public class DomainContent extends EppResource
 
   /** Loads and returns the fully qualified host names of all linked nameservers. */
   public ImmutableSortedSet<String> loadNameserverHostNames() {
-    return ofy().load()
+    return ofy()
-        .keys(getNameservers().stream().map(VKey::getOfyKey).collect(toImmutableSet())).values()
+        .load()
+        .keys(getNameservers().stream().map(VKey::getOfyKey).collect(toImmutableSet()))
+        .values()
         .stream()
         .map(HostResource::getHostName)
         .collect(toImmutableSortedSet(Ordering.natural()));

core/src/main/java/google/registry/model/domain/DomainHistory.java

@@ -38,10 +38,10 @@ import javax.persistence.JoinTable;
 @Entity
 @javax.persistence.Table(
     indexes = {
-        @javax.persistence.Index(columnList = "creationTime"),
+      @javax.persistence.Index(columnList = "creationTime"),
-        @javax.persistence.Index(columnList = "historyRegistrarId"),
+      @javax.persistence.Index(columnList = "historyRegistrarId"),
-        @javax.persistence.Index(columnList = "historyType"),
+      @javax.persistence.Index(columnList = "historyType"),
-        @javax.persistence.Index(columnList = "historyModificationTime")
+      @javax.persistence.Index(columnList = "historyModificationTime")
     })
 public class DomainHistory extends HistoryEntry {
   // Store DomainContent instead of DomainBase so we don't pick up its @Id

core/src/main/java/google/registry/model/domain/token/AllocationToken.java

@@ -29,12 +29,14 @@ import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Range;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Index;
 import com.googlecode.objectify.annotation.Mapify;
+import com.googlecode.objectify.annotation.OnLoad;
 import google.registry.flows.EppException;
 import google.registry.flows.domain.DomainFlowUtils;
 import google.registry.model.BackupGroupRoot;
@@ -108,9 +110,22 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
    */
   double discountFraction;
 
+  /** Whether the discount fraction (if any) also applies to premium names. Defaults to false. */
+  boolean discountPremiums;
+
+  /** Up to how many years of initial creation receive the discount (if any). Defaults to 1. */
+  int discountYears = 1;
+
   /** The type of the token, either single-use or unlimited-use. */
-  // TODO(b/130301183): this should not be nullable, we can remove this once we're sure it isn't
+  TokenType tokenType;
-  @Nullable TokenType tokenType;
+
+  // TODO: Remove onLoad once all allocation tokens are migrated to have a discountYears of 1.
+  @OnLoad
+  void onLoad() {
+    if (discountYears == 0) {
+      discountYears = 1;
+    }
+  }
 
   /**
    * Promotional token validity periods.
@@ -146,8 +161,8 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     return token;
   }
 
-  public Key<HistoryEntry> getRedemptionHistoryEntry() {
+  public Optional<Key<HistoryEntry>> getRedemptionHistoryEntry() {
-    return redemptionHistoryEntry;
+    return Optional.ofNullable(redemptionHistoryEntry);
   }
 
   public boolean isRedeemed() {
@@ -174,6 +189,16 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     return discountFraction;
   }
 
+  public boolean shouldDiscountPremiums() {
+    return discountPremiums;
+  }
+
+  public int getDiscountYears() {
+    // Allocation tokens created prior to the addition of the discountYears field will have a value
+    // of 0 for it, but it should be the default value of 1 to retain the previous behavior.
+    return Math.max(1, discountYears);
+  }
+
   public TokenType getTokenType() {
     return tokenType;
   }
@@ -193,6 +218,7 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
 
   /** A builder for constructing {@link AllocationToken} objects, since they are immutable. */
   public static class Builder extends Buildable.Builder<AllocationToken> {
+
     public Builder() {}
 
     private Builder(AllocationToken instance) {
@@ -210,6 +236,12 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
           getInstance().redemptionHistoryEntry == null
               || TokenType.SINGLE_USE.equals(getInstance().tokenType),
           "Redemption history entry can only be specified for SINGLE_USE tokens");
+      checkArgument(
+          getInstance().discountFraction > 0 || !getInstance().discountPremiums,
+          "Discount premiums can only be specified along with a discount fraction");
+      checkArgument(
+          getInstance().discountFraction > 0 || getInstance().discountYears == 1,
+          "Discount years can only be specified along with a discount fraction");
       if (getInstance().domainName != null) {
         try {
           DomainFlowUtils.validateDomainName(getInstance().domainName);
@@ -258,10 +290,26 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     }
 
     public Builder setDiscountFraction(double discountFraction) {
+      checkArgument(
+          Range.closed(0.0d, 1.0d).contains(discountFraction),
+          "Discount fraction must be between 0 and 1 inclusive");
       getInstance().discountFraction = discountFraction;
       return this;
     }
 
+    public Builder setDiscountPremiums(boolean discountPremiums) {
+      getInstance().discountPremiums = discountPremiums;
+      return this;
+    }
+
+    public Builder setDiscountYears(int discountYears) {
+      checkArgument(
+          Range.closed(1, 10).contains(discountYears),
+          "Discount years must be between 1 and 10 inclusive");
+      getInstance().discountYears = discountYears;
+      return this;
+    }
+
     public Builder setTokenType(TokenType tokenType) {
       checkState(getInstance().tokenType == null, "Token type can only be set once");
       getInstance().tokenType = tokenType;

core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java

@@ -115,7 +115,20 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
       description =
           "A discount off the base price for the first year between 0.0 and 1.0. Default is 0.0,"
               + " i.e. no discount.")
-  private double discountFraction;
+  private Double discountFraction;
+
+  @Parameter(
+      names = {"--discount_premiums"},
+      description =
+          "Whether the discount is valid for premium names in addition to standard ones. Default"
+              + " is false.",
+      arity = 1)
+  private Boolean discountPremiums;
+
+  @Parameter(
+      names = {"--discount_years"},
+      description = "The number of years the discount applies for. Default is 1, max value is 10.")
+  private Integer discountYears;
 
   @Parameter(
       names = "--token_status_transitions",
@@ -170,8 +183,10 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
                             .setToken(t)
                             .setTokenType(tokenType == null ? SINGLE_USE : tokenType)
                             .setAllowedClientIds(ImmutableSet.copyOf(nullToEmpty(allowedClientIds)))
-                            .setAllowedTlds(ImmutableSet.copyOf(nullToEmpty(allowedTlds)))
+                            .setAllowedTlds(ImmutableSet.copyOf(nullToEmpty(allowedTlds)));
-                            .setDiscountFraction(discountFraction);
+                    Optional.ofNullable(discountFraction).ifPresent(token::setDiscountFraction);
+                    Optional.ofNullable(discountPremiums).ifPresent(token::setDiscountPremiums);
+                    Optional.ofNullable(discountYears).ifPresent(token::setDiscountYears);
                     Optional.ofNullable(tokenStatusTransitions)
                         .ifPresent(token::setTokenStatusTransitions);
                     Optional.ofNullable(domainNames)

core/src/main/java/google/registry/tools/GetAllocationTokenCommand.java

@@ -27,7 +27,7 @@ import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.token.AllocationToken;
 import java.util.Collection;
 import java.util.List;
-import java.util.Objects;
+import java.util.Optional;
 
 /** Command to show allocation tokens. */
 @Parameters(separators = " =", commandDescription = "Show allocation token(s)")
@@ -55,11 +55,11 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
       if (loadedTokens.containsKey(token)) {
         AllocationToken loadedToken = loadedTokens.get(token);
         System.out.println(loadedToken.toString());
-        if (loadedToken.getRedemptionHistoryEntry() == null) {
+        if (!loadedToken.getRedemptionHistoryEntry().isPresent()) {
           System.out.printf("Token %s was not redeemed.\n", token);
         } else {
           DomainBase domain =
-              domains.get(loadedToken.getRedemptionHistoryEntry().<DomainBase>getParent());
+              domains.get(loadedToken.getRedemptionHistoryEntry().get().<DomainBase>getParent());
           if (domain == null) {
             System.out.printf("ERROR: Token %s was redeemed but domain can't be loaded.\n", token);
           } else {
@@ -80,7 +80,8 @@ final class GetAllocationTokenCommand implements CommandWithRemoteApi {
     ImmutableList<Key<DomainBase>> domainKeys =
         tokens.stream()
             .map(AllocationToken::getRedemptionHistoryEntry)
-            .filter(Objects::nonNull)
+            .filter(Optional::isPresent)
+            .map(Optional::get)
             .map(Key::<DomainBase>getParent)
             .collect(toImmutableList());
     ImmutableMap.Builder<Key<DomainBase>, DomainBase> domainsBuilder = new ImmutableMap.Builder<>();

core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java

@@ -66,6 +66,19 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
               + "i.e. no discount.")
   private Double discountFraction;
 
+  @Parameter(
+      names = {"--discount_premiums"},
+      description =
+          "Whether the discount is valid for premium names in addition to standard ones. Default"
+              + " is false.",
+      arity = 1)
+  private Boolean discountPremiums;
+
+  @Parameter(
+      names = {"--discount_years"},
+      description = "The number of years the discount applies for. Default is 1, max value is 10.")
+  private Integer discountYears;
+
   @Parameter(
       names = "--token_status_transitions",
       converter = TokenStatusTransitions.class,
@@ -122,6 +135,8 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     Optional.ofNullable(allowedTlds)
         .ifPresent(tlds -> builder.setAllowedTlds(ImmutableSet.copyOf(tlds)));
     Optional.ofNullable(discountFraction).ifPresent(builder::setDiscountFraction);
+    Optional.ofNullable(discountPremiums).ifPresent(builder::setDiscountPremiums);
+    Optional.ofNullable(discountYears).ifPresent(builder::setDiscountYears);
     Optional.ofNullable(tokenStatusTransitions).ifPresent(builder::setTokenStatusTransitions);
     return builder.build();
   }

core/src/test/java/google/registry/flows/domain/DomainCheckFlowTest.java

@@ -282,13 +282,14 @@ class DomainCheckFlowTest extends ResourceCheckFlowTestCase<DomainCheckFlow, Dom
   }
 
   @Test
-  void testSuccess_allocationTokenPromotion() throws Exception {
+  void testSuccess_allocationTokenPromotion_singleYear() throws Exception {
     createTld("example");
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
             .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
+            .setDiscountYears(2)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
                     .put(START_OF_TIME, TokenStatus.NOT_STARTED)
@@ -300,6 +301,69 @@ class DomainCheckFlowTest extends ResourceCheckFlowTestCase<DomainCheckFlow, Dom
     runFlowAssertResponse(loadFile("domain_check_allocationtoken_fee_response.xml"));
   }
 
+  @Test
+  void testSuccess_allocationTokenPromotion_multiYearAndPremiums() throws Exception {
+    createTld("example");
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("rich.example")
+            .setDiscountFraction(0.9)
+            .setDiscountYears(3)
+            .setDiscountPremiums(true)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().minusDays(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    setEppInput(
+        "domain_check_allocationtoken_promotion.xml", ImmutableMap.of("DOMAIN", "rich.example"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_check_allocationtoken_promotion_response.xml",
+            new ImmutableMap.Builder<String, String>()
+                .put("DOMAIN", "rich.example")
+                .put("COST_1YR", "10.00")
+                .put("COST_2YR", "20.00")
+                .put("COST_5YR", "230.00")
+                .put("FEE_CLASS", "<fee:class>premium</fee:class>")
+                .build()));
+  }
+
+  @Test
+  void testSuccess_allocationTokenPromotion_multiYear() throws Exception {
+    createTld("tld");
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("single.tld")
+            .setDiscountFraction(0.444)
+            .setDiscountYears(2)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().minusDays(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    setEppInput(
+        "domain_check_allocationtoken_promotion.xml", ImmutableMap.of("DOMAIN", "single.tld"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_check_allocationtoken_promotion_response.xml",
+            new ImmutableMap.Builder<String, String>()
+                .put("DOMAIN", "single.tld")
+                .put("COST_1YR", "7.23")
+                .put("COST_2YR", "14.46")
+                .put("COST_5YR", "53.46")
+                .put("FEE_CLASS", "")
+                .build()));
+  }
+
   @Test
   void testSuccess_promotionNotActive() throws Exception {
     createTld("example");

core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -68,6 +68,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Ordering;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig;
 import google.registry.flows.EppException;
@@ -150,12 +151,12 @@ import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
-import google.registry.model.domain.token.AllocationToken.TokenType;
 import google.registry.model.poll.PendingActionNotificationResponse.DomainPendingActionNotificationResponse;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
 import google.registry.model.registry.Registry;
+import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
@@ -435,7 +436,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Test
   void testFailure_invalidAllocationToken() {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     EppException thrown = assertThrows(InvalidAllocationTokenException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
@@ -445,7 +448,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
   void testFailure_reservedDomainCreate_allocationTokenIsForADifferentDomain() {
     // Try to register a reserved domain name with an allocation token valid for a different domain
     // name.
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld", "YEARS", "2"));
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
@@ -464,7 +468,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
   void testFailure_nonreservedDomainCreate_allocationTokenIsForADifferentDomain() {
     // Try to register a non-reserved domain name with an allocation token valid for a different
     // domain name.
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
@@ -481,7 +487,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Test
   void testFailure_alreadyRedemeedAllocationToken() {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
@@ -497,7 +505,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Test
   void testSuccess_validAllocationToken_isRedeemed() throws Exception {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     AllocationToken token =
         persistResource(
@@ -508,12 +518,14 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     HistoryEntry historyEntry =
         ofy().load().type(HistoryEntry.class).ancestor(reloadResourceByForeignKey()).first().now();
     assertThat(ofy().load().entity(token).now().getRedemptionHistoryEntry())
-        .isEqualTo(Key.create(historyEntry));
+        .hasValue(Key.create(historyEntry));
   }
 
   @Test
   void testSuccess_validAllocationToken_multiUse() throws Exception {
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     persistContactsAndHosts();
     allocationToken =
         persistResource(
@@ -531,7 +543,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     runFlow();
     assertSuccessfulCreate("tld", ImmutableSet.of(), allocationToken);
     clock.advanceOneMilli();
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "otherexample.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "otherexample.tld", "YEARS", "2"));
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "otherexample.tld")));
   }
@@ -543,8 +557,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistContactsAndHosts("foo.tld");
     assertTransactionalFlow(true);
     String expectedResponseXml =
-        loadFile(
+        loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.foo.tld"));
-            "domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.foo.tld"));
     runFlowAssertResponse(CommitMode.LIVE, UserPrivileges.NORMAL, expectedResponseXml);
     assertSuccessfulCreate("foo.tld", ImmutableSet.of());
     assertNoLordn();
@@ -576,8 +589,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         "domain_create_registration_encoded_signed_mark.xml",
         ImmutableMap.of("DOMAIN", "test-validate.tld", "PHASE", "open"));
     persistContactsAndHosts();
-    EppException thrown =
+    EppException thrown = assertThrows(SignedMarksOnlyDuringSunriseException.class, this::runFlow);
-        assertThrows(SignedMarksOnlyDuringSunriseException.class, this::runFlow);
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
@@ -1184,13 +1196,14 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         Registry.get("tld")
             .asBuilder()
             .setTldStateTransitions(
-                ImmutableSortedMap.of(
+                new ImmutableSortedMap.Builder<DateTime, TldState>(Ordering.natural())
-                    START_OF_TIME, PREDELEGATION,
+                    .put(START_OF_TIME, PREDELEGATION)
-                    DateTime.parse("1999-01-01T00:00:00Z"), QUIET_PERIOD,
+                    .put(DateTime.parse("1999-01-01T00:00:00Z"), QUIET_PERIOD)
-                    // The anchor tenant is created here, on 1999-04-03
+                    .put(DateTime.parse("1999-07-01T00:00:00Z"), START_DATE_SUNRISE)
-                    DateTime.parse("1999-07-01T00:00:00Z"), START_DATE_SUNRISE,
+                    .put(DateTime.parse("2000-01-01T00:00:00Z"), GENERAL_AVAILABILITY)
-                    DateTime.parse("2000-01-01T00:00:00Z"), GENERAL_AVAILABILITY))
+                    .build())
             .build());
+    // The anchor tenant is created during the quiet period, on 1999-04-03.
     setEppInput("domain_create_anchor_allocationtoken.xml");
     persistContactsAndHosts();
     runFlowAssertResponse(loadFile("domain_create_anchor_response.xml"));
@@ -1210,7 +1223,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                 .build());
     // Despite the domain being FULLY_BLOCKED, the non-superuser create succeeds the domain is also
     // RESERVED_FOR_SPECIFIC_USE and the correct allocation token is passed.
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld", "YEARS", "2"));
     persistContactsAndHosts();
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "resdom.tld")));
@@ -1233,7 +1247,8 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                 .setTokenType(SINGLE_USE)
                 .setDomainName("resdom.tld")
                 .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld", "YEARS", "2"));
     persistContactsAndHosts();
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "resdom.tld")));
@@ -1247,7 +1262,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         ofy().load().key(Key.create(AllocationToken.class, token)).now();
     assertThat(reloadedToken.isRedeemed()).isTrue();
     assertThat(reloadedToken.getRedemptionHistoryEntry())
-        .isEqualTo(Key.create(getHistoryEntries(reloadResourceByForeignKey()).get(0)));
+        .hasValue(Key.create(getHistoryEntries(reloadResourceByForeignKey()).get(0)));
   }
 
   private void assertAllocationTokenWasNotRedeemed(String token) {
@@ -1264,7 +1279,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
@@ -1274,7 +1289,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .build())
             .build());
     clock.advanceOneMilli();
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.tld")));
     BillingEvent.OneTime billingEvent =
@@ -1284,14 +1301,131 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
   }
 
   @Test
-  void testSuccess_promotionDoesNotApplyToPremiumPrice() {
+  void testSuccess_allocationToken_multiYearDiscount_maxesAtTokenDiscountYears() throws Exception {
-    // At the moment, discounts cannot apply to premium domains
+    // 2yrs @ $13 + 3yrs @ $13 * (1 - 0.73) = $36.53
+    runTest_allocationToken_multiYearDiscount(false, 0.73, 3, Money.of(USD, 36.53));
+  }
+
+  @Test
+  void testSuccess_allocationToken_multiYearDiscount_maxesAtNumRegistrationYears()
+      throws Exception {
+    // 5yrs @ $13 * (1 - 0.276) = $47.06
+    runTest_allocationToken_multiYearDiscount(false, 0.276, 10, Money.of(USD, 47.06));
+  }
+
+  void runTest_allocationToken_multiYearDiscount(
+      boolean discountPremiums, double discountFraction, int discountYears, Money expectedPrice)
+      throws Exception {
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("example.tld")
+            .setDiscountFraction(discountFraction)
+            .setDiscountYears(discountYears)
+            .setDiscountPremiums(discountPremiums)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().plusMillis(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusSeconds(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    clock.advanceOneMilli();
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "5"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_create_response_wildcard.xml",
+            new ImmutableMap.Builder<String, String>()
+                .put("DOMAIN", "example.tld")
+                .put("CRDATE", "1999-04-03T22:00:00.0Z")
+                .put("EXDATE", "2004-04-03T22:00:00.0Z")
+                .build()));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(ofy().load().type(BillingEvent.OneTime.class));
+    assertThat(billingEvent.getTargetId()).isEqualTo("example.tld");
+    assertThat(billingEvent.getCost()).isEqualTo(expectedPrice);
+  }
+
+  @Test
+  void testSuccess_allocationToken_multiYearDiscount_worksForPremiums() throws Exception {
     createTld("example");
     persistContactsAndHosts();
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(SINGLE_USE)
+            .setDomainName("rich.example")
+            .setDiscountFraction(0.98)
+            .setDiscountYears(2)
+            .setDiscountPremiums(true)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().plusMillis(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusSeconds(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    clock.advanceOneMilli();
+    setEppInput(
+        "domain_create_premium_allocationtoken.xml",
+        ImmutableMap.of("YEARS", "3", "FEE", "104.00"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_create_response_premium.xml",
+            ImmutableMap.of("EXDATE", "2002-04-03T22:00:00.0Z", "FEE", "104.00")));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(ofy().load().type(BillingEvent.OneTime.class));
+    assertThat(billingEvent.getTargetId()).isEqualTo("rich.example");
+    // 1yr @ $100 + 2yrs @ $100 * (1 - 0.98) = $104
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 104.00));
+  }
+
+  @Test
+  void testSuccess_allocationToken_singleYearDiscount_worksForPremiums() throws Exception {
+    createTld("example");
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(SINGLE_USE)
+            .setDomainName("rich.example")
+            .setDiscountFraction(0.95555)
+            .setDiscountPremiums(true)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(clock.nowUtc().plusMillis(1), TokenStatus.VALID)
+                    .put(clock.nowUtc().plusSeconds(1), TokenStatus.ENDED)
+                    .build())
+            .build());
+    clock.advanceOneMilli();
+    setEppInput(
+        "domain_create_premium_allocationtoken.xml",
+        ImmutableMap.of("YEARS", "3", "FEE", "204.44"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_create_response_premium.xml",
+            ImmutableMap.of("EXDATE", "2002-04-03T22:00:00.0Z", "FEE", "204.44")));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(ofy().load().type(BillingEvent.OneTime.class));
+    assertThat(billingEvent.getTargetId()).isEqualTo("rich.example");
+    // 2yrs @ $100 + 1yr @ $100 * (1 - 0.95555) = $204.44
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 204.44));
+  }
+
+  @Test
+  void testSuccess_promotionDoesNotApplyToPremiumPrice() {
+    // Discounts only apply to premium domains if the token is explicitly configured to allow it.
+    createTld("example");
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
@@ -1301,7 +1435,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .build())
             .build());
     clock.advanceOneMilli();
-    setEppInput("domain_create_premium_allocationtoken.xml");
+    setEppInput(
+        "domain_create_premium_allocationtoken.xml",
+        ImmutableMap.of("YEARS", "2", "FEE", "193.50"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenInvalidForPremiumNameException.class, this::runFlow))
         .marshalsToXml();
@@ -1313,7 +1449,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
@@ -1322,7 +1458,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .put(clock.nowUtc().plusDays(60), TokenStatus.ENDED)
                     .build())
             .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenNotInPromotionException.class, this::runFlow))
         .marshalsToXml();
@@ -1334,7 +1472,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setAllowedTlds(ImmutableSet.of("example"))
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
@@ -1344,7 +1482,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
                     .build())
             .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenNotValidForTldException.class, this::runFlow))
         .marshalsToXml();
@@ -1356,7 +1496,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     persistResource(
         new AllocationToken.Builder()
             .setToken("abc123")
-            .setTokenType(TokenType.UNLIMITED_USE)
+            .setTokenType(UNLIMITED_USE)
             .setAllowedClientIds(ImmutableSet.of("someClientId"))
             .setDiscountFraction(0.5)
             .setTokenStatusTransitions(
@@ -1366,7 +1506,9 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
                     .put(clock.nowUtc().plusDays(1), TokenStatus.ENDED)
                     .build())
             .build());
-    setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
     assertAboutEppExceptions()
         .that(assertThrows(AllocationTokenNotValidForRegistrarException.class, this::runFlow))
         .marshalsToXml();
@@ -1562,7 +1704,11 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     // Modify the Registrar to block premium names.
     persistResource(loadRegistrar("TheRegistrar").asBuilder().setBlockPremiumNames(true).build());
     runFlowAssertResponse(
-        CommitMode.LIVE, SUPERUSER, loadFile("domain_create_response_premium.xml"));
+        CommitMode.LIVE,
+        SUPERUSER,
+        loadFile(
+            "domain_create_response_premium.xml",
+            ImmutableMap.of("EXDATE", "2001-04-03T22:00:00.0Z", "FEE", "200.00")));
     assertSuccessfulCreate("example", ImmutableSet.of());
   }
 

core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java

@@ -44,7 +44,7 @@ import org.junit.jupiter.api.Test;
 class AllocationTokenTest extends EntityTestCase {
 
   @BeforeEach
-  void setup() {
+  void beforeEach() {
     createTld("foo");
   }
 
@@ -59,6 +59,8 @@ class AllocationTokenTest extends EntityTestCase {
                 .setAllowedTlds(ImmutableSet.of("dev", "app"))
                 .setAllowedClientIds(ImmutableSet.of("TheRegistrar, NewRegistrar"))
                 .setDiscountFraction(0.5)
+                .setDiscountPremiums(true)
+                .setDiscountYears(3)
                 .setTokenStatusTransitions(
                     ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
                         .put(START_OF_TIME, NOT_STARTED)
@@ -155,7 +157,7 @@ class AllocationTokenTest extends EntityTestCase {
   }
 
   @Test
-  void testBuild_invalidTLD() {
+  void testBuild_invalidTld() {
     IllegalArgumentException thrown =
         assertThrows(
             IllegalArgumentException.class,
@@ -268,6 +270,50 @@ class AllocationTokenTest extends EntityTestCase {
     assertTerminal(CANCELLED);
   }
 
+  @Test
+  void testSetDiscountFractionTooHigh() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountFraction(1.1));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount fraction must be between 0 and 1 inclusive");
+  }
+
+  @Test
+  void testSetDiscountFractionTooLow() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountFraction(-.0001));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount fraction must be between 0 and 1 inclusive");
+  }
+
+  @Test
+  void testSetDiscountYearsTooHigh() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountYears(11));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount years must be between 1 and 10 inclusive");
+  }
+
+  @Test
+  void testSetDiscountYearsTooLow() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new AllocationToken.Builder().setDiscountYears(0));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount years must be between 1 and 10 inclusive");
+  }
+
   @Test
   void testBuild_noTokenType() {
     IllegalArgumentException thrown =
@@ -295,6 +341,38 @@ class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token must not be blank");
   }
 
+  @Test
+  void testBuild_discountPremiumsRequiresDiscountFraction() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new AllocationToken.Builder()
+                    .setToken("abc")
+                    .setTokenType(SINGLE_USE)
+                    .setDiscountPremiums(true)
+                    .build());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount premiums can only be specified along with a discount fraction");
+  }
+
+  @Test
+  void testBuild_discountYearsRequiresDiscountFraction() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new AllocationToken.Builder()
+                    .setToken("abc")
+                    .setTokenType(SINGLE_USE)
+                    .setDiscountYears(2)
+                    .build());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Discount years can only be specified along with a discount fraction");
+  }
+
   private void assertBadInitialTransition(TokenStatus status) {
     assertBadTransition(
         ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()

core/src/test/java/google/registry/model/history/ContactHistoryTest.java

@@ -69,9 +69,11 @@ public class ContactHistoryTest extends EntityTestCase {
   }
 
   static void assertContactHistoriesEqual(ContactHistory one, ContactHistory two) {
-    assertAboutImmutableObjects().that(one)
+    assertAboutImmutableObjects()
+        .that(one)
         .isEqualExceptFields(two, "contactBase", "contactRepoId", "parent");
-    assertAboutImmutableObjects().that(one.getContactBase())
+    assertAboutImmutableObjects()
+        .that(one.getContactBase())
         .isEqualExceptFields(two.getContactBase(), "repoId");
   }
 }

core/src/test/java/google/registry/model/history/DomainHistoryTest.java

@@ -83,7 +83,8 @@ public class DomainHistoryTest extends EntityTestCase {
   }
 
   static void assertDomainHistoriesEqual(DomainHistory one, DomainHistory two) {
-    assertAboutImmutableObjects().that(one)
+    assertAboutImmutableObjects()
+        .that(one)
         .isEqualExceptFields(two, "domainContent", "domainRepoId", "parent");
   }
 }

core/src/test/java/google/registry/testing/DatastoreHelper.java

@@ -26,6 +26,7 @@ import static google.registry.config.RegistryConfig.getContactAndHostRoidSuffix;
 import static google.registry.config.RegistryConfig.getContactAutomaticTransferLength;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.model.ResourceTransferUtils.createTransferResponse;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
@@ -71,6 +72,7 @@ import google.registry.model.domain.DomainAuthInfo;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.eppcommon.AuthInfo.PasswordAuth;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
@@ -832,6 +834,12 @@ public class DatastoreHelper {
         .collect(onlyElement());
   }
 
+  public static void assertAllocationTokens(AllocationToken... expectedTokens) {
+    assertThat(ofy().load().type(AllocationToken.class).list())
+        .comparingElementsUsing(immutableObjectCorrespondence("updateTimestamp", "creationTime"))
+        .containsExactlyElementsIn(expectedTokens);
+  }
+
   /** Returns a newly allocated, globally unique domain repoId of the format HEX-TLD. */
   public static String generateNewDomainRoid(String tld) {
     return createDomainRepoId(ObjectifyService.allocateId(), tld);

core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
 import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatastoreHelper.assertAllocationTokens;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
@@ -32,11 +33,9 @@ import static org.mockito.Mockito.verify;
 import com.beust.jcommander.ParameterException;
 import com.google.appengine.tools.remoteapi.RemoteApiException;
 import com.google.common.base.Joiner;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Maps;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
@@ -160,6 +159,8 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         "--allowed_client_ids", "TheRegistrar,NewRegistrar",
         "--allowed_tlds", "tld,example",
         "--discount_fraction", "0.5",
+        "--discount_premiums", "true",
+        "--discount_years", "6",
         "--token_status_transitions",
             String.format(
                 "\"%s=NOT_STARTED,%s=VALID,%s=ENDED\"", START_OF_TIME, promoStart, promoEnd));
@@ -170,6 +171,8 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
             .setAllowedClientIds(ImmutableSet.of("TheRegistrar", "NewRegistrar"))
             .setAllowedTlds(ImmutableSet.of("tld", "example"))
             .setDiscountFraction(0.5)
+            .setDiscountPremiums(true)
+            .setDiscountYears(6)
             .setTokenStatusTransitions(
                 ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
                     .put(START_OF_TIME, TokenStatus.NOT_STARTED)
@@ -309,26 +312,6 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
         .isEqualTo("For UNLIMITED_USE tokens, must specify --token_status_transitions");
   }
 
-  private void assertAllocationTokens(AllocationToken... expectedTokens) {
-    // Using ImmutableObject comparison here is tricky because the creation/updated timestamps are
-    // neither easy nor valuable to test here.
-    ImmutableMap<String, AllocationToken> actualTokens =
-        Maps.uniqueIndex(ofy().load().type(AllocationToken.class), AllocationToken::getToken);
-    assertThat(actualTokens).hasSize(expectedTokens.length);
-    for (AllocationToken expectedToken : expectedTokens) {
-      AllocationToken match = actualTokens.get(expectedToken.getToken());
-      assertThat(match).isNotNull();
-      assertThat(match.getRedemptionHistoryEntry())
-          .isEqualTo(expectedToken.getRedemptionHistoryEntry());
-      assertThat(match.getAllowedClientIds()).isEqualTo(expectedToken.getAllowedClientIds());
-      assertThat(match.getAllowedTlds()).isEqualTo(expectedToken.getAllowedTlds());
-      assertThat(match.getDiscountFraction()).isEqualTo(expectedToken.getDiscountFraction());
-      assertThat(match.getTokenStatusTransitions())
-          .isEqualTo(expectedToken.getTokenStatusTransitions());
-      assertThat(match.getTokenType()).isEqualTo(expectedToken.getTokenType());
-    }
-  }
-
   private AllocationToken createToken(
       String token,
       @Nullable Key<HistoryEntry> redemptionHistoryEntry,

core/src/test/java/google/registry/tools/UpdateAllocationTokensCommandTest.java

@@ -77,6 +77,24 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(token).getDiscountFraction()).isEqualTo(0.15);
   }
 
+  @Test
+  void testUpdateDiscountPremiums() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setDiscountFraction(0.5).setDiscountPremiums(false).build());
+    runCommandForced("--prefix", "token", "--discount_premiums", "true");
+    assertThat(reloadResource(token).shouldDiscountPremiums()).isTrue();
+    runCommandForced("--prefix", "token", "--discount_premiums", "false");
+    assertThat(reloadResource(token).shouldDiscountPremiums()).isFalse();
+  }
+
+  @Test
+  void testUpdateDiscountYears() throws Exception {
+    AllocationToken token = persistResource(builderWithPromo().setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--discount_years", "4");
+    assertThat(reloadResource(token).getDiscountYears()).isEqualTo(4);
+  }
+
   @Test
   void testUpdateStatusTransitions() throws Exception {
     DateTime now = DateTime.now(UTC);

core/src/test/resources/google/registry/flows/domain/domain_check_allocationtoken_promotion.xml

@@ -0,0 +1,39 @@
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <check>
+      <domain:check
+          xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>%DOMAIN%</domain:name>
+      </domain:check>
+    </check>
+    <extension>
+      <allocationToken:allocationToken
+          xmlns:allocationToken=
+              "urn:ietf:params:xml:ns:allocationToken-1.0">
+        abc123
+      </allocationToken:allocationToken>
+      <fee:check
+          xmlns:fee="urn:ietf:params:xml:ns:fee-0.6">
+        <fee:domain>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+        </fee:domain>
+        <fee:domain>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">2</fee:period>
+        </fee:domain>
+        <fee:domain>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">5</fee:period>
+        </fee:domain>
+      </fee:check>
+    </extension>
+    <clTRID>ABC-12345</clTRID>
+  </command>
+</epp>

core/src/test/resources/google/registry/flows/domain/domain_check_allocationtoken_promotion_response.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<epp xmlns:launch="urn:ietf:params:xml:ns:launch-1.0" xmlns:secDNS="urn:ietf:params:xml:ns:secDNS-1.1" xmlns:host="urn:ietf:params:xml:ns:host-1.0" xmlns:fee11="urn:ietf:params:xml:ns:fee-0.11" xmlns:fee12="urn:ietf:params:xml:ns:fee-0.12" xmlns:fee="urn:ietf:params:xml:ns:fee-0.6" xmlns:rgp="urn:ietf:params:xml:ns:rgp-1.0" xmlns="urn:ietf:params:xml:ns:epp-1.0" xmlns:domain="urn:ietf:params:xml:ns:domain-1.0" xmlns:contact="urn:ietf:params:xml:ns:contact-1.0">
+  <response>
+    <result code="1000">
+      <msg>Command completed successfully</msg>
+    </result>
+    <resData>
+      <domain:chkData>
+        <domain:cd>
+          <domain:name avail="true">%DOMAIN%</domain:name>
+        </domain:cd>
+      </domain:chkData>
+    </resData>
+    <extension>
+      <fee:chkData>
+        <fee:cd>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+          <fee:fee description="create">%COST_1YR%</fee:fee>
+          %FEE_CLASS%
+        </fee:cd>
+        <fee:cd>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">2</fee:period>
+          <fee:fee description="create">%COST_2YR%</fee:fee>
+          %FEE_CLASS%
+        </fee:cd>
+        <fee:cd>
+          <fee:name>%DOMAIN%</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">5</fee:period>
+          <fee:fee description="create">%COST_5YR%</fee:fee>
+          %FEE_CLASS%
+        </fee:cd>
+      </fee:chkData>
+    </extension>
+    <trID>
+      <clTRID>ABC-12345</clTRID>
+      <svTRID>server-trid</svTRID>
+    </trID>
+  </response>
+</epp>

core/src/test/resources/google/registry/flows/domain/domain_create_allocationtoken.xml

@@ -4,7 +4,7 @@
       <domain:create
        xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
         <domain:name>%DOMAIN%</domain:name>
-        <domain:period unit="y">2</domain:period>
+        <domain:period unit="y">%YEARS%</domain:period>
         <domain:ns>
           <domain:hostObj>ns1.example.net</domain:hostObj>
           <domain:hostObj>ns2.example.net</domain:hostObj>

core/src/test/resources/google/registry/flows/domain/domain_create_premium_allocationtoken.xml

@@ -4,7 +4,7 @@
       <domain:create
           xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
         <domain:name>rich.example</domain:name>
-        <domain:period unit="y">2</domain:period>
+        <domain:period unit="y">%YEARS%</domain:period>
         <domain:ns>
           <domain:hostObj>ns1.example.net</domain:hostObj>
           <domain:hostObj>ns2.example.net</domain:hostObj>
@@ -25,7 +25,7 @@
       </allocationToken:allocationToken>
       <fee:create xmlns:fee="urn:ietf:params:xml:ns:fee-0.12">
         <fee:currency>USD</fee:currency>
-        <fee:fee>193.5</fee:fee>
+        <fee:fee>%FEE%</fee:fee>
       </fee:create>
     </extension>
     <clTRID>ABC-12345</clTRID>

core/src/test/resources/google/registry/flows/domain/domain_create_response_premium.xml

@@ -8,13 +8,13 @@
        xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
         <domain:name>rich.example</domain:name>
         <domain:crDate>1999-04-03T22:00:00.0Z</domain:crDate>
-        <domain:exDate>2001-04-03T22:00:00.0Z</domain:exDate>
+        <domain:exDate>%EXDATE%</domain:exDate>
       </domain:creData>
     </resData>
     <extension>
       <fee:creData xmlns:fee="urn:ietf:params:xml:ns:fee-0.12">
         <fee:currency>USD</fee:currency>
-        <fee:fee description="create">200.00</fee:fee>
+        <fee:fee description="create">%FEE%</fee:fee>
       </fee:creData>
     </extension>
     <trID>

core/src/test/resources/google/registry/flows/domain/domain_create_response_wildcard.xml

@@ -0,0 +1,19 @@
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <response>
+    <result code="1000">
+      <msg>Command completed successfully</msg>
+    </result>
+    <resData>
+      <domain:creData
+       xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>%DOMAIN%</domain:name>
+        <domain:crDate>%CRDATE%</domain:crDate>
+        <domain:exDate>%EXDATE%</domain:exDate>
+      </domain:creData>
+    </resData>
+    <trID>
+      <clTRID>ABC-12345</clTRID>
+      <svTRID>server-trid</svTRID>
+    </trID>
+  </response>
+</epp>

core/src/test/resources/google/registry/model/schema.txt

@@ -231,12 +231,14 @@ class google.registry.model.domain.secdns.DelegationSignerData {
 }
 class google.registry.model.domain.token.AllocationToken {
   @Id java.lang.String token;
+  boolean discountPremiums;
   com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> redemptionHistoryEntry;
   double discountFraction;
   google.registry.model.CreateAutoTimestamp creationTime;
   google.registry.model.UpdateAutoTimestamp updateTimestamp;
   google.registry.model.common.TimedTransitionProperty<google.registry.model.domain.token.AllocationToken$TokenStatus, google.registry.model.domain.token.AllocationToken$TokenStatusTransition> tokenStatusTransitions;
   google.registry.model.domain.token.AllocationToken$TokenType tokenType;
+  int discountYears;
   java.lang.String domainName;
   java.util.Set<java.lang.String> allowedClientIds;
   java.util.Set<java.lang.String> allowedTlds;

util/src/main/java/google/registry/util/CollectionUtils.java

@@ -75,38 +75,39 @@ public class CollectionUtils {
   }
 
   /** Defensive copy helper for {@link Set}. */
-  public static <V> ImmutableSet<V> nullSafeImmutableCopy(Set<V> data) {
+  public static <V> ImmutableSet<V> nullSafeImmutableCopy(@Nullable Set<V> data) {
     return data == null ? null : ImmutableSet.copyOf(data);
   }
 
   /** Defensive copy helper for {@link List}. */
-  public static <V> ImmutableList<V> nullSafeImmutableCopy(List<V> data) {
+  public static <V> ImmutableList<V> nullSafeImmutableCopy(@Nullable List<V> data) {
     return data == null ? null : ImmutableList.copyOf(data);
   }
 
   /** Defensive copy helper for {@link Set}. */
-  public static <V> ImmutableSet<V> nullToEmptyImmutableCopy(Set<V> data) {
+  public static <V> ImmutableSet<V> nullToEmptyImmutableCopy(@Nullable Set<V> data) {
     return data == null ? ImmutableSet.of() : ImmutableSet.copyOf(data);
   }
 
   /** Defensive copy helper for {@link Set}. */
-  public static <V extends Comparable<V>>
+  public static <V extends Comparable<V>> ImmutableSortedSet<V> nullToEmptyImmutableSortedCopy(
-      ImmutableSortedSet<V> nullToEmptyImmutableSortedCopy(Set<V> data) {
+      @Nullable Set<V> data) {
     return data == null ? ImmutableSortedSet.of() : ImmutableSortedSet.copyOf(data);
   }
 
   /** Defensive copy helper for {@link SortedMap}. */
-  public static <K, V> ImmutableSortedMap<K, V> nullToEmptyImmutableCopy(SortedMap<K, V> data) {
+  public static <K, V> ImmutableSortedMap<K, V> nullToEmptyImmutableCopy(
+      @Nullable SortedMap<K, V> data) {
     return data == null ? ImmutableSortedMap.of() : ImmutableSortedMap.copyOfSorted(data);
   }
 
   /** Defensive copy helper for {@link List}. */
-  public static <V> ImmutableList<V> nullToEmptyImmutableCopy(List<V> data) {
+  public static <V> ImmutableList<V> nullToEmptyImmutableCopy(@Nullable List<V> data) {
     return data == null ? ImmutableList.of() : ImmutableList.copyOf(data);
   }
 
   /** Defensive copy helper for {@link Map}. */
-  public static <K, V> ImmutableMap<K, V> nullToEmptyImmutableCopy(Map<K, V> data) {
+  public static <K, V> ImmutableMap<K, V> nullToEmptyImmutableCopy(@Nullable Map<K, V> data) {
     return data == null ? ImmutableMap.of() : ImmutableMap.copyOf(data);
   }
 

util/src/main/java/google/registry/util/PreconditionsUtils.java

@@ -29,33 +29,36 @@ public class PreconditionsUtils {
    * preferable to throw an IAE instead of an NPE, such as where we want an IAE to indicate that
    * it's just a bad argument/parameter and reserve NPEs for bugs and unexpected null values.
    */
-  public static <T> T checkArgumentNotNull(T reference) {
+  public static <T> T checkArgumentNotNull(@Nullable T reference) {
     checkArgument(reference != null);
     return reference;
   }
 
   /** Checks whether the provided reference is null, throws IAE if it is, and returns it if not. */
-  public static <T> T checkArgumentNotNull(T reference, @Nullable Object errorMessage) {
+  public static <T> T checkArgumentNotNull(@Nullable T reference, @Nullable Object errorMessage) {
     checkArgument(reference != null, errorMessage);
     return reference;
   }
 
   /** Checks whether the provided reference is null, throws IAE if it is, and returns it if not. */
   public static <T> T checkArgumentNotNull(
-      T reference, @Nullable String errorMessageTemplate, @Nullable Object... errorMessageArgs) {
+      @Nullable T reference,
+      @Nullable String errorMessageTemplate,
+      @Nullable Object... errorMessageArgs) {
     checkArgument(reference != null, errorMessageTemplate, errorMessageArgs);
     return reference;
   }
 
   /** Checks if the provided Optional is present, returns its value if so, and throws IAE if not. */
-  public static <T> T checkArgumentPresent(Optional<T> reference) {
+  public static <T> T checkArgumentPresent(@Nullable Optional<T> reference) {
     checkArgumentNotNull(reference);
     checkArgument(reference.isPresent());
     return reference.get();
   }
 
   /** Checks if the provided Optional is present, returns its value if so, and throws IAE if not. */
-  public static <T> T checkArgumentPresent(Optional<T> reference, @Nullable Object errorMessage) {
+  public static <T> T checkArgumentPresent(
+      @Nullable Optional<T> reference, @Nullable Object errorMessage) {
     checkArgumentNotNull(reference, errorMessage);
     checkArgument(reference.isPresent(), errorMessage);
     return reference.get();
@@ -63,7 +66,7 @@ public class PreconditionsUtils {
 
   /** Checks if the provided Optional is present, returns its value if so, and throws IAE if not. */
   public static <T> T checkArgumentPresent(
-      Optional<T> reference,
+      @Nullable Optional<T> reference,
       @Nullable String errorMessageTemplate,
       @Nullable Object... errorMessageArgs) {
     checkArgumentNotNull(reference, errorMessageTemplate, errorMessageArgs);