Use the cheapest default token when multiple are available (#2990)

Previously we would just use the first one we found. This is a valid behavior, but we want to change it so that we apply the cheapest default if multiple are available (this way we avoid having to go back after the fact and give refunds).
2026-04-10 19:47:20 +00:00 · 2026-03-27 13:19:38 -04:00
parent 3513364c97
commit a129a0dc21
16 changed files with 215 additions and 74 deletions
--- a/core/src/main/java/google/registry/bsa/persistence/BsaLabelUtils.java
+++ b/core/src/main/java/google/registry/bsa/persistence/BsaLabelUtils.java
@@ -100,8 +100,7 @@ public final class BsaLabelUtils {
    ImmutableList<VKey<BsaLabel>> queriedLabels =
        domainLabels.stream().map(BsaLabel::vKey).collect(toImmutableList());
    return cacheBsaLabels.getAll(queriedLabels).values().stream()
-        .filter(Optional::isPresent)
-        .map(Optional::get)
+        .flatMap(Optional::stream)
        .map(BsaLabel::getLabel)
        .collect(toImmutableSet());
  }
--- a/core/src/main/java/google/registry/flows/FlowReporter.java
+++ b/core/src/main/java/google/registry/flows/FlowReporter.java
@@ -100,8 +100,7 @@ public class FlowReporter {
  public static ImmutableSet<String> extractTlds(Iterable<String> domainNames) {
    return Streams.stream(domainNames)
        .map(FlowReporter::extractTld)
-        .filter(Optional::isPresent)
-        .map(Optional::get)
+        .flatMap(Optional::stream)
        .collect(toImmutableSet());
  }

--- a/core/src/main/java/google/registry/flows/domain/DomainCheckFlow.java
+++ b/core/src/main/java/google/registry/flows/domain/DomainCheckFlow.java
@@ -225,7 +225,8 @@ public final class DomainCheckFlow implements TransactionalFlow {
      ImmutableSet<InternetDomainName> bsaBlockedDomainNames,
      ImmutableMap<String, TldState> tldStates,
      ImmutableMap<String, InternetDomainName> parsedDomains,
-      DateTime now) {
+      DateTime now)
+      throws EppException {
    InternetDomainName idn = parsedDomains.get(domainName);
    Optional<AllocationToken> token;
    try {
@@ -238,7 +239,9 @@ public final class DomainCheckFlow implements TransactionalFlow {
              eppInput.getSingleExtension(AllocationTokenExtension.class),
              Tld.get(idn.parent().toString()),
              domainName,
-              FeeQueryCommandExtensionItem.CommandName.CREATE);
+              FeeQueryCommandExtensionItem.CommandName.CREATE,
+              Optional.empty(),
+              pricingLogic);
    } catch (AllocationTokenFlowUtils.NonexistentAllocationTokenException
        | AllocationTokenFlowUtils.AllocationTokenInvalidException e) {
      // The provided token was catastrophically invalid in some way
@@ -317,7 +320,9 @@ public final class DomainCheckFlow implements TransactionalFlow {
                  eppInput.getSingleExtension(AllocationTokenExtension.class),
                  tld,
                  domainName,
-                  feeCheckItem.getCommandName());
+                  feeCheckItem.getCommandName(),
+                  Optional.empty(),
+                  pricingLogic);
        } catch (AllocationTokenFlowUtils.NonexistentAllocationTokenException
            | AllocationTokenFlowUtils.AllocationTokenInvalidException e) {
          // The provided token was catastrophically invalid in some way
--- a/core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java
+++ b/core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java
@@ -264,7 +264,9 @@ public final class DomainCreateFlow implements MutatingFlow {
            eppInput.getSingleExtension(AllocationTokenExtension.class),
            tld,
            command.getDomainName(),
-            CommandName.CREATE);
+            CommandName.CREATE,
+            Optional.of(years),
+            pricingLogic);
    boolean defaultTokenUsed =
        allocationToken.map(t -> t.getTokenType().equals(TokenType.DEFAULT_PROMO)).orElse(false);
    boolean isAnchorTenant =
--- a/core/src/main/java/google/registry/flows/domain/DomainPricingLogic.java
+++ b/core/src/main/java/google/registry/flows/domain/DomainPricingLogic.java
@@ -67,7 +67,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.
   */
-  FeesAndCredits getCreatePrice(
+  public FeesAndCredits getCreatePrice(
      Tld tld,
      String domainName,
      DateTime dateTime,
@@ -193,8 +193,8 @@ public final class DomainPricingLogic {
  }

  /** Returns a new restore price for the pricer. */
-  FeesAndCredits getRestorePrice(Tld tld, String domainName, DateTime dateTime, boolean isExpired)
-      throws EppException {
+  public FeesAndCredits getRestorePrice(
+      Tld tld, String domainName, DateTime dateTime, boolean isExpired) throws EppException {
    DomainPrices domainPrices = getPricesForDomainName(domainName, dateTime);
    FeesAndCredits.Builder feesAndCredits =
        new FeesAndCredits.Builder()
@@ -216,7 +216,7 @@ public final class DomainPricingLogic {
  }

  /** Returns a new transfer price for the pricer. */
-  FeesAndCredits getTransferPrice(
+  public FeesAndCredits getTransferPrice(
      Tld tld, String domainName, DateTime dateTime, @Nullable BillingRecurrence billingRecurrence)
      throws EppException {
    FeesAndCredits renewPrice =
@@ -239,7 +239,8 @@ public final class DomainPricingLogic {
  }

  /** Returns a new update price for the pricer. */
-  FeesAndCredits getUpdatePrice(Tld tld, String domainName, DateTime dateTime) throws EppException {
+  public FeesAndCredits getUpdatePrice(Tld tld, String domainName, DateTime dateTime)
+      throws EppException {
    CurrencyUnit currency = tld.getCurrency();
    BaseFee feeOrCredit = Fee.create(zeroInCurrency(currency), FeeType.UPDATE, false);
    return customLogic.customizeUpdatePrice(
--- a/core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java
+++ b/core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java
@@ -169,6 +169,8 @@ public final class DomainRenewFlow implements MutatingFlow {
    Domain existingDomain = loadAndVerifyExistence(Domain.class, targetId, now);
    String tldStr = existingDomain.getTld();
    Tld tld = Tld.get(tldStr);
+    int years = command.getPeriod().getValue();
+
    Optional<AllocationToken> allocationToken =
        AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault(
            registrarId,
@@ -176,7 +178,9 @@ public final class DomainRenewFlow implements MutatingFlow {
            eppInput.getSingleExtension(AllocationTokenExtension.class),
            tld,
            existingDomain.getDomainName(),
-            CommandName.RENEW);
+            CommandName.RENEW,
+            Optional.of(years),
+            pricingLogic);
    boolean defaultTokenUsed =
        allocationToken
            .map(t -> t.getTokenType().equals(AllocationToken.TokenType.DEFAULT_PROMO))
@@ -186,7 +190,6 @@ public final class DomainRenewFlow implements MutatingFlow {
    // If client passed an applicable static token this updates the domain
    existingDomain = maybeApplyBulkPricingRemovalToken(existingDomain, allocationToken);

-    int years = command.getPeriod().getValue();
    DateTime newExpirationTime =
        leapSafeAddYears(existingDomain.getRegistrationExpirationTime(), years);  // Uncapped
    validateRegistrationPeriod(now, newExpirationTime);
--- a/core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java
+++ b/core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java
@@ -15,7 +15,6 @@
 package google.registry.flows.domain.token;

 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.isDomainPremium;
@@ -24,21 +23,25 @@ import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 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.flows.EppException.AssociationProhibitsOperationException;
 import google.registry.flows.EppException.AuthorizationErrorException;
 import google.registry.flows.EppException.StatusProhibitsOperationException;
-import google.registry.model.billing.BillingBase;
+import google.registry.flows.domain.DomainPricingLogic;
+import google.registry.model.billing.BillingBase.RenewalPriceBehavior;
 import google.registry.model.billing.BillingRecurrence;
 import google.registry.model.domain.Domain;
 import google.registry.model.domain.fee.FeeQueryCommandExtensionItem.CommandName;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenBehavior;
+import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.domain.token.AllocationTokenExtension;
 import google.registry.model.reporting.HistoryEntry.HistoryEntryId;
 import google.registry.model.tld.Tld;
 import google.registry.persistence.VKey;
+import java.math.BigDecimal;
 import java.util.Map;
 import java.util.Optional;
 import org.joda.time.DateTime;
@@ -91,8 +94,10 @@ public class AllocationTokenFlowUtils {
      Optional<AllocationTokenExtension> extension,
      Tld tld,
      String domainName,
-      CommandName commandName)
-      throws NonexistentAllocationTokenException, AllocationTokenInvalidException {
+      CommandName commandName,
+      Optional<Integer> years,
+      DomainPricingLogic pricingLogic)
+      throws EppException {
    Optional<AllocationToken> fromExtension =
        loadAllocationTokenFromExtension(registrarId, domainName, now, extension);
    if (fromExtension.isPresent()
@@ -100,7 +105,8 @@ public class AllocationTokenFlowUtils {
            InternetDomainName.from(domainName), fromExtension.get(), commandName, now)) {
      return fromExtension;
    }
-    return checkForDefaultToken(tld, domainName, commandName, registrarId, now);
+    return checkForDefaultToken(
+        tld, domainName, commandName, registrarId, now, years, pricingLogic);
  }

  /** Verifies that the given domain can have a bulk pricing token removed from it. */
@@ -133,7 +139,7 @@ public class AllocationTokenFlowUtils {
    BillingRecurrence newBillingRecurrence =
        tm().loadByKey(domain.getAutorenewBillingEvent())
            .asBuilder()
-            .setRenewalPriceBehavior(BillingBase.RenewalPriceBehavior.DEFAULT)
+            .setRenewalPriceBehavior(RenewalPriceBehavior.DEFAULT)
            .setRenewalPrice(null)
            .build();

@@ -182,37 +188,70 @@ public class AllocationTokenFlowUtils {
   * token found on the TLD's default token list will be returned.
   */
  private static Optional<AllocationToken> checkForDefaultToken(
-      Tld tld, String domainName, CommandName commandName, String registrarId, DateTime now) {
+      Tld tld,
+      String domainName,
+      CommandName commandName,
+      String registrarId,
+      DateTime now,
+      Optional<Integer> years,
+      DomainPricingLogic pricingLogic)
+      throws EppException {
    ImmutableList<VKey<AllocationToken>> tokensFromTld = tld.getDefaultPromoTokens();
    if (isNullOrEmpty(tokensFromTld)) {
      return Optional.empty();
    }
-    Map<VKey<AllocationToken>, Optional<AllocationToken>> tokens =
-        AllocationToken.getAll(tokensFromTld);
-    checkState(
-        !isNullOrEmpty(tokens), "Failure while loading default TLD tokens from the database");
-    // Iterate over the list to maintain token ordering (since we return the first valid token)
    ImmutableList<AllocationToken> tokenList =
-        tokensFromTld.stream()
-            .map(tokens::get)
-            .filter(Optional::isPresent)
-            .map(Optional::get)
+        AllocationToken.getAll(tokensFromTld).values().stream()
+            .flatMap(Optional::stream)
+            // Filter to tokens that are 1. valid in general 2. valid for this particular request
+            .filter(
+                token -> {
+                  try {
+                    validateTokenEntity(token, registrarId, domainName, now);
+                  } catch (AllocationTokenInvalidException e) {
+                    return false;
+                  }
+                  return tokenIsValidAgainstDomain(
+                      InternetDomainName.from(domainName), token, commandName, now);
+                })
            .collect(toImmutableList());
-
-    // Check if any of the tokens are valid for this domain registration
+    // We can't compute the costs directly in the stream due to the checked EppException
+    ImmutableMap.Builder<AllocationToken, BigDecimal> tokenCosts = new ImmutableMap.Builder<>();
    for (AllocationToken token : tokenList) {
-      try {
-        validateTokenEntity(token, registrarId, domainName, now);
-      } catch (AllocationTokenInvalidException e) {
-        // Token is not valid for this registrar, etc. -- continue trying tokens
-        continue;
-      }
-      if (tokenIsValidAgainstDomain(InternetDomainName.from(domainName), token, commandName, now)) {
-        return Optional.of(token);
-      }
+      tokenCosts.put(
+          token,
+          getSampleCostWithToken(tld, domainName, token, commandName, now, years, pricingLogic));
    }
-    // No valid default token found
-    return Optional.empty();
+    return tokenCosts.build().entrySet().stream()
+        .min(Map.Entry.comparingByValue())
+        .map(Map.Entry::getKey);
+  }
+
+  private static BigDecimal getSampleCostWithToken(
+      Tld tld,
+      String domainName,
+      AllocationToken token,
+      CommandName commandName,
+      DateTime now,
+      Optional<Integer> years,
+      DomainPricingLogic pricingLogic)
+      throws EppException {
+    int yearsForAction = years.orElse(1);
+    // We only support token discounts on creates or renews
+    return switch (commandName) {
+      case CREATE ->
+          pricingLogic
+              .getCreatePrice(
+                  tld, domainName, now, yearsForAction, false, false, Optional.of(token))
+              .getTotalCost()
+              .getAmount();
+      case RENEW ->
+          pricingLogic
+              .getRenewPrice(tld, domainName, now, yearsForAction, null, Optional.of(token))
+              .getTotalCost()
+              .getAmount();
+      default -> BigDecimal.ZERO;
+    };
  }

  /** Loads a given token and validates it against the registrar, time, etc */
@@ -253,8 +292,7 @@ public class AllocationTokenFlowUtils {
    // Tokens without status transitions will just have a single-entry NOT_STARTED map, so only
    // check the status transitions map if it's non-trivial.
    if (token.getTokenStatusTransitions().size() > 1
-        && !AllocationToken.TokenStatus.VALID.equals(
-            token.getTokenStatusTransitions().getValueAtTime(now))) {
+        && !TokenStatus.VALID.equals(token.getTokenStatusTransitions().getValueAtTime(now))) {
      throw new AllocationTokenNotInPromotionException();
    }

@@ -303,7 +341,7 @@ public class AllocationTokenFlowUtils {
    AllocationTokenNotValidForDomainException() {
      super("Alloc token invalid for domain");
    }
-    }
+  }

  /** The allocation token is invalid. */
  public static class NonexistentAllocationTokenException extends AuthorizationErrorException {
--- a/core/src/main/java/google/registry/model/tld/label/ReservedList.java
+++ b/core/src/main/java/google/registry/model/tld/label/ReservedList.java
@@ -261,8 +261,7 @@ public final class ReservedList
  public static ImmutableSet<ReservedList> loadReservedLists(
      ImmutableSet<String> reservedListNames) {
    return cache.getAll(reservedListNames).values().stream()
-        .filter(Optional::isPresent)
-        .map(Optional::get)
+        .flatMap(Optional::stream)
        .collect(toImmutableSet());
  }

--- a/core/src/main/java/google/registry/rdap/RdapJsonFormatter.java
+++ b/core/src/main/java/google/registry/rdap/RdapJsonFormatter.java
@@ -579,8 +579,7 @@ public class RdapJsonFormatter {
      ImmutableList<RdapRegistrarPocEntity> registrarPocs =
          registrar.getPocsFromReplica().stream()
              .map(RdapJsonFormatter::makeRdapJsonForRegistrarPoc)
-              .filter(Optional::isPresent)
-              .map(Optional::get)
+              .flatMap(Optional::stream)
              .filter(
                  poc ->
                      outputDataType == OutputDataType.FULL
--- a/core/src/main/java/google/registry/tools/GetAllocationTokenCommand.java
+++ b/core/src/main/java/google/registry/tools/GetAllocationTokenCommand.java
@@ -89,8 +89,7 @@ final class GetAllocationTokenCommand implements Command {
    ImmutableList<VKey<Domain>> domainKeys =
        tokens.stream()
            .map(AllocationToken::getRedemptionHistoryId)
-            .filter(Optional::isPresent)
-            .map(Optional::get)
+            .flatMap(Optional::stream)
            .map(hi -> tm().loadByKey(VKey.create(DomainHistory.class, hi)))
            .map(dh -> VKey.create(Domain.class, dh.getRepoId()))
            .collect(toImmutableList());
--- a/core/src/main/java/google/registry/tools/server/ListPremiumListsAction.java
+++ b/core/src/main/java/google/registry/tools/server/ListPremiumListsAction.java
@@ -54,8 +54,7 @@ public final class ListPremiumListsAction extends ListObjectsAction<PremiumList>
                tm().loadAllOf(PremiumList.class).stream()
                    .map(PremiumList::getName)
                    .map(PremiumListDao::getLatestRevision)
-                    .filter(Optional::isPresent)
-                    .map(Optional::get)
+                    .flatMap(Optional::stream)
                    .collect(toImmutableSortedSet(Comparator.comparing(PremiumList::getName))));
  }
 }
--- a/core/src/main/java/google/registry/tools/server/ListReservedListsAction.java
+++ b/core/src/main/java/google/registry/tools/server/ListReservedListsAction.java
@@ -52,8 +52,7 @@ public final class ListReservedListsAction extends ListObjectsAction<ReservedLis
                tm().loadAllOf(ReservedList.class).stream()
                    .map(ReservedList::getName)
                    .map(ReservedListDao::getLatestRevision)
-                    .filter(Optional::isPresent)
-                    .map(Optional::get)
+                    .flatMap(Optional::stream)
                    .collect(toImmutableSortedSet(Comparator.comparing(ReservedList::getName))));
  }
 }
--- a/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java
@@ -1577,7 +1577,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
    persistHosts();
    setupDefaultToken("aaaaa", 0, "TheRegistrar");
    setupDefaultTokenWithDiscount();
-    runTest_defaultToken("aaaaa");
+    runTest_defaultToken("bbbbb");
  }

  @Test
--- a/core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java
@@ -1243,7 +1243,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, Domain>
            new AllocationToken.Builder()
                .setToken("aaaaa")
                .setTokenType(DEFAULT_PROMO)
-                .setDiscountFraction(0.5)
+                .setDiscountFraction(0.9)
                .setDiscountYears(1)
                .setAllowedTlds(ImmutableSet.of("tld"))
                .build());
@@ -1271,8 +1271,8 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, Domain>
    assertThat(billingEvent.getTargetId()).isEqualTo("example.tld");
    assertThat(billingEvent.getAllocationToken().get().getKey())
        .isEqualTo(defaultToken1.getToken());
-    // Price is 50% off the first year only. Non-discounted price is $11.
-    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 16.5));
+    // Price is 90% off the first year only. Non-discounted price is $11.
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 12.10));
  }

  @Test
@@ -1412,7 +1412,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, Domain>
  }

  @Test
-  void testSuccess_onlyUsesFirstValidToken() throws Exception {
+  void testSuccess_usesCheapestValidToken() throws Exception {
    setEppInput("domain_renew.xml", ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
    persistDomain();
    AllocationToken defaultToken1 =
@@ -1459,10 +1459,9 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, Domain>
    BillingEvent billingEvent =
        Iterables.getOnlyElement(DatabaseHelper.loadAllOf(BillingEvent.class));
    assertThat(billingEvent.getTargetId()).isEqualTo("example.tld");
-    assertThat(billingEvent.getAllocationToken().get().getKey())
-        .isEqualTo(defaultToken2.getToken());
-    // Price is 50% off the first year only. Non-discounted price is $11.
-    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 16.5));
+    assertThat(billingEvent.getAllocationToken().get().getKey()).isEqualTo("ccccc");
+    // Price is 75% off the first year only. Non-discounted price is $11.
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, 13.75));
  }

  @Test
--- a/core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java
+++ b/core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java
@@ -35,6 +35,8 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.net.InternetDomainName;
 import google.registry.flows.EppException;
+import google.registry.flows.custom.DomainPricingCustomLogic;
+import google.registry.flows.domain.DomainPricingLogic;
 import google.registry.flows.domain.token.AllocationTokenFlowUtils.AllocationTokenNotInPromotionException;
 import google.registry.flows.domain.token.AllocationTokenFlowUtils.AllocationTokenNotValidForRegistrarException;
 import google.registry.flows.domain.token.AllocationTokenFlowUtils.NonexistentAllocationTokenException;
@@ -65,6 +67,9 @@ class AllocationTokenFlowUtilsTest {
  private final AllocationTokenExtension allocationTokenExtension =
      mock(AllocationTokenExtension.class);

+  private final DomainPricingLogic domainPricingLogic =
+      new DomainPricingLogic(new DomainPricingCustomLogic(null, null, null));
+
  private Tld tld;

  @BeforeEach
@@ -140,7 +145,9 @@ class AllocationTokenFlowUtilsTest {
                Optional.of(allocationTokenExtension),
                tld,
                "example.tld",
-                CommandName.CREATE))
+                CommandName.CREATE,
+                Optional.of(1),
+                domainPricingLogic))
        .hasValue(token);
  }

@@ -154,7 +161,9 @@ class AllocationTokenFlowUtilsTest {
                Optional.empty(),
                tld,
                "example.tld",
-                CommandName.CREATE))
+                CommandName.CREATE,
+                Optional.of(1),
+                domainPricingLogic))
        .hasValue(defaultToken);
  }

@@ -176,7 +185,9 @@ class AllocationTokenFlowUtilsTest {
                Optional.of(allocationTokenExtension),
                tld,
                "example.tld",
-                CommandName.CREATE))
+                CommandName.CREATE,
+                Optional.of(1),
+                domainPricingLogic))
        .hasValue(defaultToken);
  }

@@ -299,7 +310,9 @@ class AllocationTokenFlowUtilsTest {
                Optional.of(allocationTokenExtension),
                tld,
                "example.tld",
-                CommandName.CREATE));
+                CommandName.CREATE,
+                Optional.of(1),
+                domainPricingLogic));
  }

  @Test
@@ -311,7 +324,9 @@ class AllocationTokenFlowUtilsTest {
                Optional.empty(),
                tld,
                "example.tld",
-                CommandName.CREATE))
+                CommandName.CREATE,
+                Optional.of(1),
+                domainPricingLogic))
        .isEmpty();
  }

@@ -329,7 +344,93 @@ class AllocationTokenFlowUtilsTest {
                Optional.of(allocationTokenExtension),
                tld,
                "example.tld",
-                CommandName.CREATE));
+                CommandName.CREATE,
+                Optional.of(1),
+                domainPricingLogic));
+  }
+
+  @Test
+  void testSuccess_default_cheaperTokenUsed() throws Exception {
+    AllocationToken cheaperToken =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("cheaperToken")
+                .setDiscountFraction(0.5)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setTokenType(DEFAULT_PROMO)
+                .build());
+    AllocationToken moreExpensiveToken =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("moreExpensiveToken")
+                .setDiscountFraction(0.1)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setTokenType(DEFAULT_PROMO)
+                .build());
+    // List the more expensive token first to ensure that we don't just pick the first valid one
+    tld =
+        persistResource(
+            tld.asBuilder()
+                .setDefaultPromoTokens(
+                    ImmutableList.of(moreExpensiveToken.createVKey(), cheaperToken.createVKey()))
+                .build());
+
+    assertThat(
+            AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault(
+                "TheRegistrar",
+                clock.nowUtc(),
+                Optional.empty(),
+                tld,
+                "example.tld",
+                CommandName.CREATE,
+                Optional.of(1),
+                domainPricingLogic))
+        .hasValue(cheaperToken);
+  }
+
+  @Test
+  void testSuccess_default_twoYearsIsCheaper() throws Exception {
+    AllocationToken longerToken =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("longerToken")
+                .setDiscountFraction(0.4)
+                .setDiscountYears(2)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setTokenType(DEFAULT_PROMO)
+                .build());
+    AllocationToken shorterToken =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("shorterToken")
+                .setDiscountFraction(0.5)
+                .setDiscountYears(1)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setTokenType(DEFAULT_PROMO)
+                .build());
+    tld =
+        persistResource(
+            tld.asBuilder()
+                .setDefaultPromoTokens(
+                    ImmutableList.of(shorterToken.createVKey(), longerToken.createVKey()))
+                .build());
+
+    // The token with the smaller discount fraction should be chosen because it runs for 2 years
+    assertThat(
+            AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault(
+                "TheRegistrar",
+                clock.nowUtc(),
+                Optional.empty(),
+                tld,
+                "example.tld",
+                CommandName.CREATE,
+                Optional.of(2),
+                domainPricingLogic))
+        .hasValue(longerToken);
  }

  private AllocationToken persistDefaultToken() {
--- a/db/src/test/java/google/registry/sql/flyway/FlywayDeadlockTest.java
+++ b/db/src/test/java/google/registry/sql/flyway/FlywayDeadlockTest.java
@@ -178,8 +178,7 @@ public class FlywayDeadlockTest {
                  .filter(line -> !line.isBlank())
                  .collect(joining(" ")))
          .map(FlywayDeadlockTest::getDdlLockedElementName)
-          .filter(Optional::isPresent)
-          .map(Optional::get)
+          .flatMap(Optional::stream)
          .collect(toImmutableSet());
    } catch (IOException e) {
      throw new RuntimeException(e);