Use a potential discount in the AllocationToken when determining domain create price

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=245458027

java/google/registry/flows/domain/DomainCheckFlow.java

@@ -44,6 +44,7 @@ import google.registry.flows.annotations.ReportingSpec;
 import google.registry.flows.custom.DomainCheckFlowCustomLogic;
 import google.registry.flows.custom.DomainCheckFlowCustomLogic.BeforeResponseParameters;
 import google.registry.flows.custom.DomainCheckFlowCustomLogic.BeforeResponseReturnData;
+import google.registry.flows.domain.token.AllocationTokenDomainCheckResults;
 import google.registry.flows.domain.token.AllocationTokenFlowUtils;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.DomainCommand.Check;
@@ -51,6 +52,7 @@ import google.registry.model.domain.fee.FeeCheckCommandExtension;
 import google.registry.model.domain.fee.FeeCheckCommandExtensionItem;
 import google.registry.model.domain.fee.FeeCheckResponseExtensionItem;
 import google.registry.model.domain.launch.LaunchCheckExtension;
+import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationTokenExtension;
 import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppinput.ResourceCommand;
@@ -150,27 +152,38 @@ public final class DomainCheckFlow implements Flow {
     Set<String> existingIds = checkResourcesExist(DomainBase.class, targetIds, now);
     Optional<AllocationTokenExtension> allocationTokenExtension =
         eppInput.getSingleExtension(AllocationTokenExtension.class);
-    ImmutableMap<InternetDomainName, String> tokenCheckResults =
+    Optional<AllocationTokenDomainCheckResults> tokenDomainCheckResults =
-        allocationTokenExtension.isPresent()
+        allocationTokenExtension.map(
-            ? allocationTokenFlowUtils.checkDomainsWithToken(
+            tokenExtension ->
+                allocationTokenFlowUtils.checkDomainsWithToken(
                     ImmutableList.copyOf(domainNames.values()),
-                allocationTokenExtension.get().getAllocationToken(),
+                    tokenExtension.getAllocationToken(),
                     clientId,
-                now)
+                    now));
-            : ImmutableMap.of();
+
     ImmutableList.Builder<DomainCheck> checks = new ImmutableList.Builder<>();
     ImmutableMap<String, TldState> tldStates =
         Maps.toMap(seenTlds, tld -> Registry.get(tld).getTldState(now));
+    ImmutableMap<InternetDomainName, String> domainCheckResults =
+        tokenDomainCheckResults
+            .map(AllocationTokenDomainCheckResults::domainCheckResults)
+            .orElse(ImmutableMap.of());
     for (String targetId : targetIds) {
       Optional<String> message =
-          getMessageForCheck(domainNames.get(targetId), existingIds, tokenCheckResults, tldStates);
+          getMessageForCheck(
+              domainNames.get(targetId),
+              existingIds,
+              domainCheckResults,
+              tldStates);
       checks.add(DomainCheck.create(!message.isPresent(), targetId, message.orElse(null)));
     }
+    Optional<AllocationToken> allocationToken =
+        tokenDomainCheckResults.flatMap(AllocationTokenDomainCheckResults::token);
     BeforeResponseReturnData responseData =
         flowCustomLogic.beforeResponse(
             BeforeResponseParameters.newBuilder()
                 .setDomainChecks(checks.build())
-                .setResponseExtensions(getResponseExtensions(domainNames, now))
+                .setResponseExtensions(getResponseExtensions(domainNames, now, allocationToken))
                 .setAsOfDate(now)
                 .build());
     return responseBuilder
@@ -199,7 +212,10 @@ public final class DomainCheckFlow implements Flow {
 
   /** Handle the fee check extension. */
   private ImmutableList<? extends ResponseExtension> getResponseExtensions(
-      ImmutableMap<String, InternetDomainName> domainNames, DateTime now) throws EppException {
+      ImmutableMap<String, InternetDomainName> domainNames,
+      DateTime now,
+      Optional<AllocationToken> allocationToken)
+      throws EppException {
     Optional<FeeCheckCommandExtension> feeCheckOpt =
         eppInput.getSingleExtension(FeeCheckCommandExtension.class);
     if (!feeCheckOpt.isPresent()) {
@@ -217,7 +233,8 @@ public final class DomainCheckFlow implements Flow {
             domainNames.get(domainName),
             feeCheck.getCurrency(),
             now,
-            pricingLogic);
+            pricingLogic,
+            allocationToken);
         responseItems.add(builder.setDomainNameIfSupported(domainName).build());
       }
     }

java/google/registry/flows/domain/DomainCreateFlow.java

@@ -277,7 +277,8 @@ public class DomainCreateFlow implements TransactionalFlow {
     Optional<FeeCreateCommandExtension> feeCreate =
         eppInput.getSingleExtension(FeeCreateCommandExtension.class);
     FeesAndCredits feesAndCredits =
-        pricingLogic.getCreatePrice(registry, targetId, now, years, isAnchorTenant);
+        pricingLogic.getCreatePrice(
+            registry, targetId, now, years, isAnchorTenant, allocationToken);
     validateFeeChallenge(targetId, now, feeCreate, feesAndCredits);
     Optional<SecDnsCreateExtension> secDnsCreate =
         validateSecDnsExtension(eppInput.getSingleExtension(SecDnsCreateExtension.class));
@@ -444,7 +445,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         eppInput.getSingleExtension(AllocationTokenExtension.class);
     return Optional.ofNullable(
         extension.isPresent()
-            ? allocationTokenFlowUtils.verifyToken(
+            ? allocationTokenFlowUtils.loadAndVerifyToken(
                 command, extension.get().getAllocationToken(), registry, clientId, now)
             : null);
   }

java/google/registry/flows/domain/DomainFlowUtils.java

@@ -550,7 +550,8 @@ public class DomainFlowUtils {
       InternetDomainName domain,
       @Nullable CurrencyUnit topLevelCurrency,
       DateTime currentDate,
-      DomainPricingLogic pricingLogic)
+      DomainPricingLogic pricingLogic,
+      Optional<AllocationToken> allocationToken)
       throws EppException {
     DateTime now = currentDate;
     // Use the custom effective date specified in the fee check request, if there is one.
@@ -588,10 +589,10 @@ public class DomainFlowUtils {
           builder.setReasonIfSupported("reserved");
         } else {
           builder.setAvailIfSupported(true);
-          // TODO(b/117145844): Once allocation token support for domain check flow is implemented,
-          // we should be able to calculate the correct price here.
           fees =
-              pricingLogic.getCreatePrice(registry, domainNameString, now, years, false).getFees();
+              pricingLogic
+                  .getCreatePrice(registry, domainNameString, now, years, false, allocationToken)
+                  .getFees();
         }
         break;
       case RENEW:

java/google/registry/flows/domain/DomainInfoFlow.java

@@ -163,7 +163,8 @@ public final class DomainInfoFlow implements Flow {
           InternetDomainName.from(targetId),
           null,
           now,
-          pricingLogic);
+          pricingLogic,
+          Optional.empty());
       extensions.add(builder.build());
     }
     return extensions.build();

java/google/registry/flows/domain/DomainPricingLogic.java

@@ -14,7 +14,7 @@
 
 package google.registry.flows.domain;
 
-import static google.registry.pricing.PricingEngineProxy.getDomainCreateCost;
+import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.pricing.PricingEngineProxy.getDomainFeeClass;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
 
@@ -30,8 +30,12 @@ import google.registry.flows.custom.DomainPricingCustomLogic.UpdatePriceParamete
 import google.registry.model.domain.fee.BaseFee;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Fee;
+import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.pricing.PremiumPricingEngine.DomainPrices;
 import google.registry.model.registry.Registry;
+import google.registry.pricing.PricingEngineProxy;
 import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.money.CurrencyUnit;
@@ -51,16 +55,27 @@ public final class DomainPricingLogic {
   @Inject
   DomainPricingLogic() {}
 
-  /** Returns a new create price for the pricer. */
+  /**
+   * Returns a new create price for the pricer.
+   *
+   * <p>If {@code allocationToken} is present and the domain is non-premium, that discount will be
+   * applied to the first year.
+   */
   public FeesAndCredits getCreatePrice(
-      Registry registry, String domainName, DateTime date, int years, boolean isAnchorTenant)
+      Registry registry,
+      String domainName,
+      DateTime date,
+      int years,
+      boolean isAnchorTenant,
+      Optional<AllocationToken> allocationToken)
       throws EppException {
     CurrencyUnit currency = registry.getCurrency();
-
+    // Domain create cost is always zero for anchor tenants
-    // Get the vanilla create cost, or 0 for anchor tenants.
+    Money domainCreateCost =
-    BigDecimal domainCreateCost =
+        isAnchorTenant
-        isAnchorTenant ? BigDecimal.ZERO : getDomainCreateCost(domainName, date, years).getAmount();
+            ? Money.of(currency, BigDecimal.ZERO)
-    BaseFee createFeeOrCredit = Fee.create(domainCreateCost, FeeType.CREATE);
+            : getDomainCreateCostWithDiscount(domainName, date, years, allocationToken);
+    BaseFee createFeeOrCredit = Fee.create(domainCreateCost.getAmount(), FeeType.CREATE);
 
     // Create fees for the cost and the EAP fee, if any.
     Fee eapFee = registry.getEapFeeFor(date);
@@ -80,7 +95,6 @@ public final class DomainPricingLogic {
             .setAsOfDate(date)
             .setYears(years)
             .build());
-
   }
 
   /** Returns a new renew price for the pricer. */
@@ -154,7 +168,7 @@ public final class DomainPricingLogic {
             .setFeesAndCredits(
                 new FeesAndCredits.Builder()
                     .setCurrency(currency)
-                    .addFeeOrCredit(feeOrCredit)
+                    .setFeesAndCredits(feeOrCredit)
                     .build())
             .setRegistry(registry)
             .setDomainName(InternetDomainName.from(domainName))
@@ -166,4 +180,24 @@ public final class DomainPricingLogic {
   public Optional<String> getFeeClass(String domainName, DateTime date) {
     return getDomainFeeClass(domainName, date);
   }
+
+  private Money getDomainCreateCostWithDiscount(
+      String domainName, DateTime date, int years, Optional<AllocationToken> allocationToken) {
+    DomainPrices domainPrices = PricingEngineProxy.getPricesForDomainName(domainName, date);
+    checkArgument(
+        !allocationToken.isPresent()
+            || allocationToken.get().getDiscountFraction() == 0.0
+            || !domainPrices.isPremium(),
+        "A nonzero discount code cannot be applied to premium domains");
+    Money oneYearCreateCost = domainPrices.getCreateCost();
+    Money totalDomainCreateCost = oneYearCreateCost.multipliedBy(years);
+    // If a discount is applicable, apply it only to the first year
+    if (allocationToken.isPresent()) {
+      Money discount =
+          oneYearCreateCost.multipliedBy(
+              allocationToken.get().getDiscountFraction(), RoundingMode.HALF_UP);
+      totalDomainCreateCost = totalDomainCreateCost.minus(discount);
+    }
+    return totalDomainCreateCost;
+  }
 }

java/google/registry/flows/domain/token/AllocationTokenCustomLogic.java

@@ -14,12 +14,14 @@
 
 package google.registry.flows.domain.token;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.net.InternetDomainName;
 import google.registry.flows.EppException;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.registry.Registry;
+import java.util.function.Function;
 import org.joda.time.DateTime;
 
 /**
@@ -43,11 +45,12 @@ public class AllocationTokenCustomLogic {
 
   /** Performs additional custom logic for performing domain checks using a token. */
   public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
-      ImmutableMap<InternetDomainName, String> checkResults,
+      ImmutableList<InternetDomainName> domainNames,
       AllocationToken token,
       String clientId,
       DateTime now) {
     // Do nothing.
-    return checkResults;
+    return domainNames.stream()
+        .collect(ImmutableMap.toImmutableMap(Function.identity(), ignored -> ""));
   }
 }

java/google/registry/flows/domain/token/AllocationTokenDomainCheckResults.java

@@ -0,0 +1,36 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// 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 google.registry.flows.domain.token;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.net.InternetDomainName;
+import google.registry.model.domain.token.AllocationToken;
+import java.util.Optional;
+
+/** Value class to represent the result of loading a token and checking domains with it. */
+@AutoValue
+public abstract class AllocationTokenDomainCheckResults {
+
+  public abstract Optional<AllocationToken> token();
+
+  public abstract ImmutableMap<InternetDomainName, String> domainCheckResults();
+
+  public static AllocationTokenDomainCheckResults create(
+      Optional<AllocationToken> allocationToken,
+      ImmutableMap<InternetDomainName, String> domainCheckResults) {
+    return new AutoValue_AllocationTokenDomainCheckResults(allocationToken, domainCheckResults);
+  }
+}

java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java

@@ -17,7 +17,9 @@ package google.registry.flows.domain.token;
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import com.google.common.net.InternetDomainName;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
@@ -29,7 +31,7 @@ import google.registry.model.domain.token.AllocationToken.TokenType;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import java.util.List;
-import java.util.function.Function;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 
@@ -49,16 +51,10 @@ public class AllocationTokenFlowUtils {
    * @return the loaded {@link AllocationToken} for that string.
    * @throws InvalidAllocationTokenException if the token doesn't exist.
    */
-  public AllocationToken verifyToken(
+  public AllocationToken loadAndVerifyToken(
       DomainCommand.Create command, String token, Registry registry, String clientId, DateTime now)
       throws EppException {
-    AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
+    AllocationToken tokenEntity = loadToken(token);
-    if (tokenEntity == null) {
-      throw new InvalidAllocationTokenException();
-    }
-    if (tokenEntity.isRedeemed()) {
-      throw new AlreadyRedeemedAllocationTokenException();
-    }
     return tokenCustomLogic.verifyToken(command, tokenEntity, registry, clientId, now);
   }
 
@@ -69,26 +65,21 @@ public class AllocationTokenFlowUtils {
    *     for a a given domain then it does not validate with this allocation token; domains that do
    *     validate have blank messages (i.e. no error).
    */
-  public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
+  public AllocationTokenDomainCheckResults checkDomainsWithToken(
       List<InternetDomainName> domainNames, String token, String clientId, DateTime now) {
-    AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
+    try {
-    String globalResult;
+      AllocationToken tokenEntity = loadToken(token);
-    if (tokenEntity == null) {
-      globalResult = new InvalidAllocationTokenException().getMessage();
-    } else if (tokenEntity.isRedeemed()) {
-      globalResult = AlreadyRedeemedAllocationTokenException.ERROR_MSG_SHORT;
-    } else {
-      globalResult = "";
-    }
-    ImmutableMap<InternetDomainName, String> checkResults =
-        domainNames
-            .stream()
-            .collect(ImmutableMap.toImmutableMap(Function.identity(), domainName -> globalResult));
       // Only call custom logic if there wasn't a global allocation token error that applies to all
       // check results. The custom logic can only add errors, not override existing errors.
-    return globalResult.isEmpty()
+      return AllocationTokenDomainCheckResults.create(
-        ? tokenCustomLogic.checkDomainsWithToken(checkResults, tokenEntity, clientId, now)
+          Optional.of(tokenEntity),
-        : checkResults;
+          tokenCustomLogic.checkDomainsWithToken(
+              ImmutableList.copyOf(domainNames), tokenEntity, clientId, now));
+    } catch (EppException e) {
+      return AllocationTokenDomainCheckResults.create(
+          Optional.empty(),
+          ImmutableMap.copyOf(Maps.toMap(domainNames, ignored -> e.getMessage())));
+    }
   }
 
   /**
@@ -102,17 +93,22 @@ public class AllocationTokenFlowUtils {
     return token.asBuilder().setRedemptionHistoryEntry(redemptionHistoryEntry).build();
   }
 
+  private AllocationToken loadToken(String token) throws EppException {
+    AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
+    if (tokenEntity == null) {
+      throw new InvalidAllocationTokenException();
+    }
+    if (tokenEntity.isRedeemed()) {
+      throw new AlreadyRedeemedAllocationTokenException();
+    }
+    return tokenEntity;
+  }
+
   /** The allocation token was already redeemed. */
   public static class AlreadyRedeemedAllocationTokenException
       extends AssociationProhibitsOperationException {
-
-    public static final String ERROR_MSG_LONG = "The allocation token was already redeemed";
-
-    /** A short error message fitting within 32 characters for use in domain check responses. */
-    public static final String ERROR_MSG_SHORT = "Alloc token was already redeemed";
-
     public AlreadyRedeemedAllocationTokenException() {
-      super(ERROR_MSG_LONG);
+      super("Alloc token was already redeemed");
     }
   }
 

javatests/google/registry/flows/domain/DomainCheckFlowTest.java

@@ -15,6 +15,7 @@
 package google.registry.flows.domain;
 
 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.eppoutput.CheckData.DomainCheck.create;
 import static google.registry.model.registry.Registry.TldState.PREDELEGATION;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
@@ -63,6 +64,7 @@ import google.registry.flows.domain.DomainFlowUtils.UnknownFeeCommandException;
 import google.registry.flows.exceptions.TooManyResourceChecksException;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.label.ReservedList;
@@ -182,6 +184,24 @@ public class DomainCheckFlowTest
         create(false, "premiumcollision.tld", "Cannot be delegated"));
   }
 
+  @Test
+  public void testSuccess_allocationTokenPromotion() throws Exception {
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(UNLIMITED_USE)
+            .setDiscountFraction(0.5)
+            .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());
+    setEppInput("domain_check_allocationtoken_fee.xml");
+    runFlowAssertResponse(loadFile("domain_check_allocationtoken_fee_response.xml"));
+  }
+
   @Test
   public void testSuccess_oneReservedInSunrise() throws Exception {
     createTld("tld", START_DATE_SUNRISE);

javatests/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -65,6 +65,7 @@ import com.google.common.collect.ImmutableList;
 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.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig;
 import google.registry.flows.EppException;
@@ -140,6 +141,8 @@ import google.registry.model.domain.launch.LaunchNotice;
 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;
@@ -151,6 +154,7 @@ import google.registry.model.reporting.DomainTransactionRecord.TransactionReport
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.monitoring.whitebox.EppMetric;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
+import java.math.BigDecimal;
 import java.util.Map;
 import javax.annotation.Nullable;
 import org.joda.money.Money;
@@ -1101,6 +1105,57 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
         .isEqualTo(Key.create(getHistoryEntries(reloadResourceByForeignKey()).get(0)));
   }
 
+  @Test
+  public void testSuccess_allocationTokenPromotion() throws Exception {
+    // A discount of 0.5 means that the first-year cost (13) is cut in half, so a discount of 6.5
+    // Note: we're asking to register it for two years so the total cost should be 13 + (13/2)
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(TokenType.UNLIMITED_USE)
+            .setDiscountFraction(0.5)
+            .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"));
+    runFlowAssertResponse(
+        loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.tld")));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(ofy().load().type(BillingEvent.OneTime.class));
+    assertThat(billingEvent.getTargetId()).isEqualTo("example.tld");
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, BigDecimal.valueOf(19.5)));
+  }
+
+  @Test
+  public void testSuccess_promotionDoesNotApplyToPremiumPrice() {
+    // At the moment, discounts cannot apply to premium domains
+    createTld("example");
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(TokenType.UNLIMITED_USE)
+            .setDiscountFraction(0.5)
+            .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");
+    assertThat(assertThrows(IllegalArgumentException.class, this::runFlow))
+        .hasMessageThat()
+        .isEqualTo("A nonzero discount code cannot be applied to premium domains");
+  }
+
   @Test
   public void testSuccess_superuserReserved() throws Exception {
     setEppInput("domain_create_reserved.xml");

javatests/google/registry/flows/domain/testdata/domain_check_allocationtoken_fee.xml

@@ -0,0 +1,41 @@
+<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>example1.tld</domain:name>
+        <domain:name>example2.tld</domain:name>
+        <domain:name>reserved.tld</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>example1.tld</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>example2.tld</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>reserved.tld</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+        </fee:domain>
+      </fee:check>
+    </extension>
+    <clTRID>ABC-12345</clTRID>
+  </command>
+</epp>

javatests/google/registry/flows/domain/testdata/domain_check_allocationtoken_fee_response.xml

@@ -0,0 +1,51 @@
+<?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">example1.tld</domain:name>
+        </domain:cd>
+        <domain:cd>
+          <domain:name avail="true">example2.tld</domain:name>
+        </domain:cd>
+        <domain:cd>
+          <domain:name avail="false">reserved.tld</domain:name>
+          <domain:reason>Reserved</domain:reason>
+        </domain:cd>
+      </domain:chkData>
+    </resData>
+    <extension>
+      <fee:chkData>
+        <fee:cd>
+          <fee:name>example2.tld</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+          <fee:fee description="create">6.50</fee:fee>
+        </fee:cd>
+        <fee:cd>
+          <fee:name>example1.tld</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+          <fee:fee description="create">6.50</fee:fee>
+        </fee:cd>
+        <fee:cd>
+          <fee:name>reserved.tld</fee:name>
+          <fee:currency>USD</fee:currency>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+          <fee:class>reserved</fee:class>
+        </fee:cd>
+      </fee:chkData>
+    </extension>
+    <trID>
+      <clTRID>ABC-12345</clTRID>
+      <svTRID>server-trid</svTRID>
+    </trID>
+  </response>
+</epp>

javatests/google/registry/flows/domain/testdata/domain_create_premium_allocationtoken.xml

@@ -0,0 +1,33 @@
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <create>
+      <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:ns>
+          <domain:hostObj>ns1.example.net</domain:hostObj>
+          <domain:hostObj>ns2.example.net</domain:hostObj>
+        </domain:ns>
+        <domain:registrant>jd1234</domain:registrant>
+        <domain:contact type="admin">sh8013</domain:contact>
+        <domain:contact type="tech">sh8013</domain:contact>
+        <domain:authInfo>
+          <domain:pw>2fooBAR</domain:pw>
+        </domain:authInfo>
+      </domain:create>
+    </create>
+    <extension>
+      <allocationToken:allocationToken
+          xmlns:allocationToken=
+              "urn:ietf:params:xml:ns:allocationToken-1.0">
+        abc123
+      </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:create>
+    </extension>
+    <clTRID>ABC-12345</clTRID>
+  </command>
+</epp>

javatests/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java

@@ -36,7 +36,7 @@ import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.testing.AppEngineRule;
 import google.registry.testing.ShardableTestCase;
-import java.util.Map;
+import java.util.function.Function;
 import org.joda.time.DateTime;
 import org.junit.Before;
 import org.junit.Rule;
@@ -63,7 +63,7 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
     AllocationTokenFlowUtils flowUtils =
         new AllocationTokenFlowUtils(new AllocationTokenCustomLogic());
     assertThat(
-            flowUtils.verifyToken(
+            flowUtils.loadAndVerifyToken(
                 createCommand("blah.tld"),
                 "tokeN",
                 Registry.get("tld"),
@@ -80,7 +80,7 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
         assertThrows(
             InvalidAllocationTokenException.class,
             () ->
-                flowUtils.verifyToken(
+                flowUtils.loadAndVerifyToken(
                     createCommand("blah.tld"),
                     "tokeN",
                     Registry.get("tld"),
@@ -99,7 +99,7 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
         assertThrows(
             IllegalStateException.class,
             () ->
-                flowUtils.verifyToken(
+                flowUtils.loadAndVerifyToken(
                     createCommand("blah.tld"),
                     "tokeN",
                     Registry.get("tld"),
@@ -115,12 +115,14 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
     AllocationTokenFlowUtils flowUtils =
         new AllocationTokenFlowUtils(new AllocationTokenCustomLogic());
     assertThat(
-            flowUtils.checkDomainsWithToken(
+            flowUtils
+                .checkDomainsWithToken(
                     ImmutableList.of(
                         InternetDomainName.from("blah.tld"), InternetDomainName.from("blah2.tld")),
                     "tokeN",
                     "TheRegistrar",
-                DateTime.now(UTC)))
+                    DateTime.now(UTC))
+                .domainCheckResults())
         .containsExactlyEntriesIn(
             ImmutableMap.of(
                 InternetDomainName.from("blah.tld"), "", InternetDomainName.from("blah2.tld"), ""))
@@ -138,12 +140,14 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
     AllocationTokenFlowUtils flowUtils =
         new AllocationTokenFlowUtils(new AllocationTokenCustomLogic());
     assertThat(
-            flowUtils.checkDomainsWithToken(
+            flowUtils
+                .checkDomainsWithToken(
                     ImmutableList.of(
                         InternetDomainName.from("blah.tld"), InternetDomainName.from("blah2.tld")),
                     "tokeN",
                     "TheRegistrar",
-                DateTime.now(UTC)))
+                    DateTime.now(UTC))
+                .domainCheckResults())
         .containsExactlyEntriesIn(
             ImmutableMap.of(
                 InternetDomainName.from("blah.tld"),
@@ -179,12 +183,14 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
     AllocationTokenFlowUtils flowUtils =
         new AllocationTokenFlowUtils(new CustomResultAllocationTokenCustomLogic());
     assertThat(
-            flowUtils.checkDomainsWithToken(
+            flowUtils
+                .checkDomainsWithToken(
                     ImmutableList.of(
                         InternetDomainName.from("blah.tld"), InternetDomainName.from("bunny.tld")),
                     "tokeN",
                     "TheRegistrar",
-                DateTime.now(UTC)))
+                    DateTime.now(UTC))
+                .domainCheckResults())
         .containsExactlyEntriesIn(
             ImmutableMap.of(
                 InternetDomainName.from("blah.tld"),
@@ -215,7 +221,7 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
 
     @Override
     public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
-        ImmutableMap<InternetDomainName, String> checkResults,
+        ImmutableList<InternetDomainName> domainNames,
         AllocationToken tokenEntity,
         String clientId,
         DateTime now) {
@@ -228,18 +234,15 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
 
     @Override
     public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
-        ImmutableMap<InternetDomainName, String> checkResults,
+        ImmutableList<InternetDomainName> domainNames,
         AllocationToken tokenEntity,
         String clientId,
         DateTime now) {
-      return checkResults
+      return domainNames.stream()
-          .entrySet()
-          .stream()
           .collect(
               ImmutableMap.toImmutableMap(
-                  Map.Entry::getKey,
+                  Function.identity(),
-                  entry ->
+                  domainName -> domainName.toString().contains("bunny") ? "fufu" : ""));
-                      entry.getKey().toString().contains("bunny") ? "fufu" : entry.getValue()));
     }
   }
 }