Add DateTime as extensibility param for allocation token logic

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

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

@@ -155,10 +155,13 @@ public final class DomainCheckFlow implements Flow {
     Set<String> existingIds = checkResourcesExist(DomainResource.class, targetIds, now);
     Optional<AllocationTokenExtension> allocationTokenExtension =
         eppInput.getSingleExtension(AllocationTokenExtension.class);
-    ImmutableMap<String, String> tokenCheckResults =
+    ImmutableMap<InternetDomainName, String> tokenCheckResults =
         allocationTokenExtension.isPresent()
             ? allocationTokenFlowUtils.checkDomainsWithToken(
-                targetIds, allocationTokenExtension.get().getAllocationToken(), clientId)
+                ImmutableList.copyOf(domainNames.values()),
+                allocationTokenExtension.get().getAllocationToken(),
+                clientId,
+                now)
             : ImmutableMap.of();
     ImmutableList.Builder<DomainCheck> checks = new ImmutableList.Builder<>();
     for (String targetId : targetIds) {
@@ -182,7 +185,7 @@ public final class DomainCheckFlow implements Flow {
   private Optional<String> getMessageForCheck(
       InternetDomainName domainName,
       Set<String> existingIds,
-      ImmutableMap<String, String> tokenCheckResults,
+      ImmutableMap<InternetDomainName, String> tokenCheckResults,
       DateTime now) {
     if (existingIds.contains(domainName.toString())) {
       return Optional.of("In use");
@@ -204,7 +207,7 @@ public final class DomainCheckFlow implements Flow {
     if (!reservationTypes.isEmpty()) {
       return Optional.of(getTypeOfHighestSeverity(reservationTypes).getMessageForCheck());
     }
-    return Optional.ofNullable(emptyToNull(tokenCheckResults.get(domainName.toString())));
+    return Optional.ofNullable(emptyToNull(tokenCheckResults.get(domainName)));
   }
 
   /** Handle the fee check extension. */

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

@@ -74,6 +74,7 @@ import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.AllocationToken;
 import google.registry.model.domain.DomainApplication;
+import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.DomainCommand.Create;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.GracePeriod;
@@ -204,7 +205,7 @@ public class DomainCreateFlow implements TransactionalFlow {
     validateClientIsLoggedIn(clientId);
     verifyRegistrarIsActive(clientId);
     DateTime now = ofy().getTransactionTime();
-    Create command = cloneAndLinkReferences((Create) resourceCommand, now);
+    DomainCommand.Create command = cloneAndLinkReferences((Create) resourceCommand, now);
     Period period = command.getPeriod();
     verifyUnitIsYears(period);
     int years = period.getValue();
@@ -261,7 +262,7 @@ public class DomainCreateFlow implements TransactionalFlow {
       }
     }
     Optional<AllocationToken> allocationToken =
-        verifyAllocationTokenIfPresent(domainName, registry, clientId);
+        verifyAllocationTokenIfPresent(command, registry, clientId, now);
     flowCustomLogic.afterValidation(
         DomainCreateFlowCustomLogic.AfterValidationParameters.newBuilder()
             .setDomainName(domainName)
@@ -384,14 +385,14 @@ public class DomainCreateFlow implements TransactionalFlow {
 
   /** Verifies and returns the allocation token if one is specified, otherwise does nothing. */
   private Optional<AllocationToken> verifyAllocationTokenIfPresent(
-      InternetDomainName domainName, Registry registry, String clientId)
+      DomainCommand.Create command, Registry registry, String clientId, DateTime now)
       throws EppException {
     Optional<AllocationTokenExtension> extension =
         eppInput.getSingleExtension(AllocationTokenExtension.class);
     return Optional.ofNullable(
         extension.isPresent()
             ? allocationTokenFlowUtils.verifyToken(
-                domainName, extension.get().getAllocationToken(), registry, clientId)
+                command, extension.get().getAllocationToken(), registry, clientId, now)
             : null);
   }
 

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

@@ -18,7 +18,9 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.net.InternetDomainName;
 import google.registry.flows.EppException;
 import google.registry.model.domain.AllocationToken;
+import google.registry.model.domain.DomainCommand;
 import google.registry.model.registry.Registry;
+import org.joda.time.DateTime;
 
 /**
  * A no-op base class for allocation token custom logic.
@@ -29,15 +31,22 @@ public class AllocationTokenCustomLogic {
 
   /** Performs additional custom logic for verifying a token. */
   public AllocationToken verifyToken(
-      InternetDomainName domainName, AllocationToken token, Registry registry, String clientId)
+      DomainCommand.Create command,
+      AllocationToken token,
+      Registry registry,
+      String clientId,
+      DateTime now)
       throws EppException {
     // Do nothing.
     return token;
   }
 
   /** Performs additional custom logic for performing domain checks using a token. */
-  public ImmutableMap<String, String> checkDomainsWithToken(
+  public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
-      ImmutableMap<String, String> checkResults, AllocationToken tokenEntity, String clientId) {
+      ImmutableMap<InternetDomainName, String> checkResults,
+      AllocationToken token,
+      String clientId,
+      DateTime now) {
     // Do nothing.
     return checkResults;
   }

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

@@ -23,14 +23,15 @@ import google.registry.flows.EppException;
 import google.registry.flows.EppException.AssociationProhibitsOperationException;
 import google.registry.flows.EppException.ParameterValueSyntaxErrorException;
 import google.registry.model.domain.AllocationToken;
+import google.registry.model.domain.DomainCommand;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import java.util.List;
 import java.util.function.Function;
 import javax.inject.Inject;
+import org.joda.time.DateTime;
 
 /** Utility functions for dealing with {@link AllocationToken}s in domain flows. */
-// TODO: Add a test class.
 public class AllocationTokenFlowUtils {
 
   final AllocationTokenCustomLogic tokenCustomLogic;
@@ -47,7 +48,7 @@ public class AllocationTokenFlowUtils {
    * @throws InvalidAllocationTokenException if the token doesn't exist.
    */
   public AllocationToken verifyToken(
-      InternetDomainName domainName, String token, Registry registry, String clientId)
+      DomainCommand.Create command, String token, Registry registry, String clientId, DateTime now)
       throws EppException {
     AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
     if (tokenEntity == null) {
@@ -56,7 +57,7 @@ public class AllocationTokenFlowUtils {
     if (tokenEntity.isRedeemed()) {
       throw new AlreadyRedeemedAllocationTokenException();
     }
-    return tokenCustomLogic.verifyToken(domainName, tokenEntity, registry, clientId);
+    return tokenCustomLogic.verifyToken(command, tokenEntity, registry, clientId, now);
   }
 
   /**
@@ -66,22 +67,26 @@ 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<String, String> checkDomainsWithToken(
+  public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
-      List<String> domainNames, String token, String clientId) {
+      List<InternetDomainName> domainNames, String token, String clientId, DateTime now) {
     AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
-    String result;
+    String globalResult;
     if (tokenEntity == null) {
-      result = new InvalidAllocationTokenException().getMessage();
+      globalResult = new InvalidAllocationTokenException().getMessage();
     } else if (tokenEntity.isRedeemed()) {
-      result = AlreadyRedeemedAllocationTokenException.ERROR_MSG_SHORT;
+      globalResult = AlreadyRedeemedAllocationTokenException.ERROR_MSG_SHORT;
     } else {
-      result = "";
+      globalResult = "";
     }
-    ImmutableMap<String, String> checkResults =
+    ImmutableMap<InternetDomainName, String> checkResults =
         domainNames
             .stream()
-            .collect(ImmutableMap.toImmutableMap(Function.identity(), domainName -> result));
+            .collect(ImmutableMap.toImmutableMap(Function.identity(), domainName -> globalResult));
-    return tokenCustomLogic.checkDomainsWithToken(checkResults, tokenEntity, clientId);
+    // 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()
+        ? tokenCustomLogic.checkDomainsWithToken(checkResults, tokenEntity, clientId, now)
+        : checkResults;
   }
 
   /** Redeems an {@link AllocationToken}, returning the redeemed copy. */

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

@@ -19,6 +19,9 @@ import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions;
 import static google.registry.testing.JUnitBackports.expectThrows;
+import static org.joda.time.DateTimeZone.UTC;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -27,10 +30,13 @@ import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.domain.token.AllocationTokenFlowUtils.InvalidAllocationTokenException;
 import google.registry.model.domain.AllocationToken;
+import google.registry.model.domain.DomainCommand;
 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 org.joda.time.DateTime;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -56,7 +62,11 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
         new AllocationTokenFlowUtils(new AllocationTokenCustomLogic());
     assertThat(
             flowUtils.verifyToken(
-                InternetDomainName.from("blah.tld"), "tokeN", Registry.get("tld"), "TheRegistrar"))
+                createCommand("blah.tld"),
+                "tokeN",
+                Registry.get("tld"),
+                "TheRegistrar",
+                DateTime.now(UTC)))
         .isEqualTo(token);
   }
 
@@ -69,10 +79,11 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
             InvalidAllocationTokenException.class,
             () ->
                 flowUtils.verifyToken(
-                    InternetDomainName.from("blah.tld"),
+                    createCommand("blah.tld"),
                     "tokeN",
                     Registry.get("tld"),
-                    "TheRegistrar"));
+                    "TheRegistrar",
+                    DateTime.now(UTC)));
     assertAboutEppExceptions().that(thrown).marshalsToXml();
   }
 
@@ -86,10 +97,11 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
             IllegalStateException.class,
             () ->
                 flowUtils.verifyToken(
-                    InternetDomainName.from("blah.tld"),
+                    createCommand("blah.tld"),
                     "tokeN",
                     Registry.get("tld"),
-                    "TheRegistrar"));
+                    "TheRegistrar",
+                    DateTime.now(UTC)));
     assertThat(thrown).hasMessageThat().isEqualTo("failed for tests");
   }
 
@@ -99,9 +111,16 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
     AllocationTokenFlowUtils flowUtils =
         new AllocationTokenFlowUtils(new AllocationTokenCustomLogic());
     assertThat(
-        flowUtils.checkDomainsWithToken(
+            flowUtils.checkDomainsWithToken(
-            ImmutableList.of("blah.tld", "blah2.tld"), "tokeN", "TheRegistrar"))
+                ImmutableList.of(
-        .containsExactlyEntriesIn(ImmutableMap.of("blah.tld", "", "blah2.tld", ""));
+                    InternetDomainName.from("blah.tld"), InternetDomainName.from("blah2.tld")),
+                "tokeN",
+                "TheRegistrar",
+                DateTime.now(UTC)))
+        .containsExactlyEntriesIn(
+            ImmutableMap.of(
+                InternetDomainName.from("blah.tld"), "", InternetDomainName.from("blah2.tld"), ""))
+        .inOrder();
   }
 
   @Test
@@ -115,13 +134,18 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
         new AllocationTokenFlowUtils(new AllocationTokenCustomLogic());
     assertThat(
             flowUtils.checkDomainsWithToken(
-                ImmutableList.of("blah.tld", "blah2.tld"), "tokeN", "TheRegistrar"))
+                ImmutableList.of(
+                    InternetDomainName.from("blah.tld"), InternetDomainName.from("blah2.tld")),
+                "tokeN",
+                "TheRegistrar",
+                DateTime.now(UTC)))
         .containsExactlyEntriesIn(
             ImmutableMap.of(
-                "blah.tld",
+                InternetDomainName.from("blah.tld"),
                 "Alloc token was already redeemed",
-                "blah2.tld",
+                InternetDomainName.from("blah2.tld"),
-                "Alloc token was already redeemed"));
+                "Alloc token was already redeemed"))
+        .inOrder();
   }
 
   @Test
@@ -134,24 +158,82 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
             IllegalStateException.class,
             () ->
                 flowUtils.checkDomainsWithToken(
-                    ImmutableList.of("blah.tld", "blah2.tld"), "tokeN", "TheRegistrar"));
+                    ImmutableList.of(
+                        InternetDomainName.from("blah.tld"), InternetDomainName.from("blah2.tld")),
+                    "tokeN",
+                    "TheRegistrar",
+                    DateTime.now(UTC)));
     assertThat(thrown).hasMessageThat().isEqualTo("failed for tests");
   }
 
+  @Test
+  public void test_checkDomainsWithToken_resultsFromCustomLogicAreIntegrated() throws Exception {
+    persistResource(new AllocationToken.Builder().setToken("tokeN").build());
+    AllocationTokenFlowUtils flowUtils =
+        new AllocationTokenFlowUtils(new CustomResultAllocationTokenCustomLogic());
+    assertThat(
+            flowUtils.checkDomainsWithToken(
+                ImmutableList.of(
+                    InternetDomainName.from("blah.tld"), InternetDomainName.from("bunny.tld")),
+                "tokeN",
+                "TheRegistrar",
+                DateTime.now(UTC)))
+        .containsExactlyEntriesIn(
+            ImmutableMap.of(
+                InternetDomainName.from("blah.tld"),
+                "",
+                InternetDomainName.from("bunny.tld"),
+                "fufu"))
+        .inOrder();
+  }
+
+  private static DomainCommand.Create createCommand(String domainName) {
+    DomainCommand.Create command = mock(DomainCommand.Create.class);
+    when(command.getFullyQualifiedDomainName()).thenReturn(domainName);
+    return command;
+  }
+
   /** An {@link AllocationTokenCustomLogic} class that throws exceptions on every method. */
   private static class FailingAllocationTokenCustomLogic extends AllocationTokenCustomLogic {
 
     @Override
     public AllocationToken verifyToken(
-        InternetDomainName domainName, AllocationToken token, Registry registry, String clientId)
+        DomainCommand.Create command,
+        AllocationToken token,
+        Registry registry,
+        String clientId,
+        DateTime now)
         throws EppException {
       throw new IllegalStateException("failed for tests");
     }
 
     @Override
-    public ImmutableMap<String, String> checkDomainsWithToken(
+    public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
-        ImmutableMap<String, String> checkResults, AllocationToken tokenEntity, String clientId) {
+        ImmutableMap<InternetDomainName, String> checkResults,
+        AllocationToken tokenEntity,
+        String clientId,
+        DateTime now) {
       throw new IllegalStateException("failed for tests");
     }
   }
+
+  /** An {@link AllocationTokenCustomLogic} class that returns custom check results for bunnies. */
+  private static class CustomResultAllocationTokenCustomLogic extends AllocationTokenCustomLogic {
+
+    @Override
+    public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
+        ImmutableMap<InternetDomainName, String> checkResults,
+        AllocationToken tokenEntity,
+        String clientId,
+        DateTime now) {
+      return checkResults
+          .entrySet()
+          .stream()
+          .collect(
+              ImmutableMap.toImmutableMap(
+                  Map.Entry::getKey,
+                  entry ->
+                      entry.getKey().toString().contains("bunny") ? "fufu" : entry.getValue()));
+    }
+  }
 }