Add promotional fields in GenerateAllocationTokensCommand

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

java/google/registry/model/domain/token/AllocationToken.java

@@ -182,6 +182,10 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     return tokenType;
   }
 
+  public TimedTransitionProperty<TokenStatus, TokenStatusTransition> getTokenStatusTransitions() {
+    return tokenStatusTransitions;
+  }
+
   @Override
   public Builder asBuilder() {
     return new Builder(clone(this));

java/google/registry/tools/GenerateAllocationTokensCommand.java

@@ -19,7 +19,9 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Queues.newArrayDeque;
 import static com.google.common.collect.Sets.difference;
 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.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.StringGenerator.DEFAULT_PASSWORD_LENGTH;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
@@ -29,10 +31,16 @@ import com.google.appengine.tools.remoteapi.RemoteApiException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 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.tools.params.TransitionListParameter.TokenStatusTransitions;
+import google.registry.util.CollectionUtils;
 import google.registry.util.NonFinalForTesting;
 import google.registry.util.Retrier;
 import google.registry.util.StringGenerator;
@@ -40,8 +48,11 @@ import java.io.File;
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Deque;
+import java.util.List;
+import java.util.Optional;
 import javax.inject.Inject;
 import javax.inject.Named;
+import org.joda.time.DateTime;
 
 /** Command to generate and persist {@link AllocationToken}s. */
 @Parameters(
@@ -54,28 +65,56 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
 
   @Parameter(
       names = {"-p", "--prefix"},
-    description = "Allocation token prefix; defaults to blank"
+      description = "Allocation token prefix; defaults to blank")
-  )
   private String prefix = "";
 
   @Parameter(
       names = {"-n", "--number"},
-    description = "The number of tokens to generate"
+      description = "The number of tokens to generate")
-  )
   private long numTokens;
 
   @Parameter(
       names = {"-d", "--domain_names_file"},
-      description = "A file with a list of newline-delimited domain names to create tokens for"
+      description = "A file with a list of newline-delimited domain names to create tokens for")
-  )
   private String domainNamesFile;
 
   @Parameter(
       names = {"-l", "--length"},
-    description = "The length of each token, exclusive of the prefix (if specified); defaults to 16"
+      description =
-  )
+          "The length of each token, exclusive of the prefix (if specified); defaults to 16")
   private int tokenLength = DEFAULT_PASSWORD_LENGTH;
 
+  @Parameter(
+      names = {"-t", "--type"},
+      description = "Type of type token, either SINGLE_USE (default) or UNLIMITED_USE")
+  private TokenType tokenType;
+
+  @Parameter(
+      names = {"--allowed_client_ids"},
+      description = "Comma-separated list of allowed client IDs, or null if all are allowed")
+  private List<String> allowedClientIds;
+
+  @Parameter(
+      names = {"--allowed_tlds"},
+      description = "Comma-separated list of allowed TLDs, or null if all are allowed")
+  private List<String> allowedTlds;
+
+  @Parameter(
+      names = {"--discount_fraction"},
+      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;
+
+  @Parameter(
+      names = "--token_status_transitions",
+      converter = TokenStatusTransitions.class,
+      validateWith = TokenStatusTransitions.class,
+      description =
+          "Comma-delimited list of token status transitions effective on specific dates, of the"
+              + " form <time>=<status>[,<time>=<status>]* where each status represents the status.")
+  private ImmutableSortedMap<DateTime, TokenStatus> tokenStatusTransitions;
+
   @Parameter(
       names = {"--dry_run"},
       description = "Do not actually persist the tokens; defaults to false")
@@ -96,6 +135,20 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
         (numTokens > 0) ^ (domainNamesFile != null),
         "Must specify either --number or --domain_names_file, but not both");
 
+    checkArgument(
+        !(UNLIMITED_USE.equals(tokenType) && CollectionUtils.isNullOrEmpty(tokenStatusTransitions)),
+        "For UNLIMITED_USE tokens, must specify --token_status_transitions");
+
+    // A list consisting solely of the empty string means user error when formatting parameters
+    checkArgument(
+        !ImmutableList.of("").equals(allowedClientIds),
+        "Either omit --allowed_client_ids if all registrars are allowed, or include a"
+            + " comma-separated list");
+
+    checkArgument(
+        !ImmutableList.of("").equals(allowedTlds),
+        "Either omit --allowed_tlds if all TLDs are allowed, or include a comma-separated list");
+
     Deque<String> domainNames;
     if (domainNamesFile == null) {
       domainNames = null;
@@ -112,17 +165,21 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
     int tokensSaved = 0;
     do {
       ImmutableSet<AllocationToken> tokens =
-          generateTokens(BATCH_SIZE)
+          generateTokens(BATCH_SIZE).stream()
-              .stream()
               .limit(numTokens - tokensSaved)
               .map(
                   t -> {
-                    AllocationToken.Builder token = new AllocationToken.Builder().setToken(t);
+                    AllocationToken.Builder token =
-                    // TODO(b/129471448): allow this to be unlimited-use as well
+                        new AllocationToken.Builder()
-                    token.setTokenType(SINGLE_USE);
+                            .setToken(t)
-                    if (domainNames != null) {
+                            .setTokenType(tokenType == null ? SINGLE_USE : tokenType)
-                      token.setDomainName(domainNames.removeFirst());
+                            .setAllowedClientIds(ImmutableSet.copyOf(nullToEmpty(allowedClientIds)))
-                    }
+                            .setAllowedTlds(ImmutableSet.copyOf(nullToEmpty(allowedTlds)))
+                            .setDiscountFraction(discountFraction);
+                    Optional.ofNullable(tokenStatusTransitions)
+                        .ifPresent(token::setTokenStatusTransitions);
+                    Optional.ofNullable(domainNames)
+                        .ifPresent(d -> token.setDomainName(d.removeFirst()));
                     return token.build();
                   })
               .collect(toImmutableSet());
@@ -149,22 +206,15 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
    */
   private ImmutableSet<String> generateTokens(int count) {
     ImmutableSet<String> candidates =
-        stringGenerator
+        stringGenerator.createStrings(tokenLength, count).stream()
-            .createStrings(tokenLength, count)
-            .stream()
             .map(s -> prefix + s)
             .collect(toImmutableSet());
     ImmutableSet<Key<AllocationToken>> existingTokenKeys =
-        candidates
+        candidates.stream()
-            .stream()
             .map(input -> Key.create(AllocationToken.class, input))
             .collect(toImmutableSet());
     ImmutableSet<String> existingTokenStrings =
-        ofy()
+        ofy().load().keys(existingTokenKeys).values().stream()
-            .load()
-            .keys(existingTokenKeys)
-            .values()
-            .stream()
             .map(AllocationToken::getToken)
             .collect(toImmutableSet());
     return ImmutableSet.copyOf(difference(candidates, existingTokenStrings));

java/google/registry/tools/params/TransitionListParameter.java

@@ -19,6 +19,7 @@ import static com.google.common.base.Preconditions.checkArgument;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Ordering;
+import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.registry.Registry.TldState;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
@@ -63,4 +64,12 @@ public abstract class TransitionListParameter<V> extends KeyValueMapParameter<Da
       return MONEY_CONVERTER.convert(value);
     }
   }
+
+  /** Converter-validator for token status transitions. */
+  public static class TokenStatusTransitions extends TransitionListParameter<TokenStatus> {
+    @Override
+    protected TokenStatus parseValue(String value) {
+      return TokenStatus.valueOf(value);
+    }
+  }
 }

javatests/google/registry/tools/GenerateAllocationTokensCommandTest.java

@@ -16,18 +16,25 @@ package google.registry.tools;
 
 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.persistResource;
 import static google.registry.testing.JUnitBackports.assertThrows;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.joda.time.DateTimeZone.UTC;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.spy;
 
+import com.beust.jcommander.ParameterException;
 import com.google.appengine.tools.remoteapi.RemoteApiException;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.testing.DeterministicStringGenerator;
 import google.registry.testing.DeterministicStringGenerator.Rule;
@@ -135,6 +142,36 @@ public class GenerateAllocationTokensCommandTest
         "foo1.tld,123456789ABCDEFG\nboo2.tld,HJKLMNPQRSTUVWXY\nbaz9.tld,Zabcdefghijkmnop");
   }
 
+  @Test
+  public void testSuccess_promotionToken() throws Exception {
+    DateTime promoStart = DateTime.now(UTC);
+    DateTime promoEnd = promoStart.plusMonths(1);
+    runCommand(
+        "--number", "1",
+        "--prefix", "promo",
+        "--type", "UNLIMITED_USE",
+        "--allowed_client_ids", "TheRegistrar,NewRegistrar",
+        "--allowed_tlds", "tld,example",
+        "--discount_fraction", "0.5",
+        "--token_status_transitions",
+            String.format(
+                "\"%s=NOT_STARTED,%s=VALID,%s=ENDED\"", START_OF_TIME, promoStart, promoEnd));
+    assertAllocationTokens(
+        new AllocationToken.Builder()
+            .setToken("promo123456789ABCDEFG")
+            .setTokenType(UNLIMITED_USE)
+            .setAllowedClientIds(ImmutableSet.of("TheRegistrar", "NewRegistrar"))
+            .setAllowedTlds(ImmutableSet.of("tld", "example"))
+            .setDiscountFraction(0.5)
+            .setTokenStatusTransitions(
+                ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
+                    .put(START_OF_TIME, TokenStatus.NOT_STARTED)
+                    .put(promoStart, TokenStatus.VALID)
+                    .put(promoEnd, TokenStatus.ENDED)
+                    .build())
+            .build());
+  }
+
   @Test
   public void testFailure_mustSpecifyNumberOfTokensOrDomainsFile() {
     IllegalArgumentException thrown =
@@ -159,15 +196,47 @@ public class GenerateAllocationTokensCommandTest
         .isEqualTo("Must specify either --number or --domain_names_file, but not both");
   }
 
+  @Test
+  public void testFailure_invalidTokenType() {
+    ParameterException thrown =
+        assertThrows(
+            ParameterException.class,
+            () -> runCommand("--number", "999", "--type", "INVALID_TYPE"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Invalid value for -t parameter. Allowed values:[SINGLE_USE, UNLIMITED_USE]");
+  }
+
+  @Test
+  public void testFailure_invalidTokenStatusTransition() {
+    assertThat(
+            assertThrows(
+                ParameterException.class,
+                () ->
+                    runCommand(
+                        "--number",
+                        "999",
+                        String.format(
+                            "--token_status_transitions=\"%s=INVALID_STATUS\"", START_OF_TIME))))
+        .hasCauseThat()
+        .isInstanceOf(IllegalArgumentException.class);
+  }
+
+  @Test
+  public void testFailure_unlimitedUseMustHaveTransitions() {
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> runCommand("--number", "999", "--type", "UNLIMITED_USE")))
+        .hasMessageThat()
+        .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 =
-        ofy()
+        ofy().load().type(AllocationToken.class).list().stream()
-            .load()
-            .type(AllocationToken.class)
-            .list()
-            .stream()
             .collect(ImmutableMap.toImmutableMap(AllocationToken::getToken, Function.identity()));
     assertThat(actualTokens).hasSize(expectedTokens.length);
     for (AllocationToken expectedToken : expectedTokens) {
@@ -175,6 +244,12 @@ public class GenerateAllocationTokensCommandTest
       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());
     }
   }