Add default tokens to TLD using nomulus tool (#1888)

* Add defualt tokens to TLD using nomulus tool

* add test

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

@@ -15,7 +15,9 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static google.registry.tools.UpdateOrDeleteAllocationTokensCommand.getTokenKeys;
 import static google.registry.util.CollectionUtils.findDuplicates;
+import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.DomainNameUtils.canonicalizeHostname;
 
 import com.beust.jcommander.Parameter;
@@ -232,6 +234,16 @@ abstract class CreateOrUpdateTldCommand extends MutatingCommand {
   )
   Integer numDnsPublishShards;
 
+  @Nullable
+  @Parameter(
+      names = "--default_tokens",
+      description =
+          "A comma-separated list of default allocation tokens to be applied to the TLD. The"
+              + " ordering of this list will determine which token is used in the case where"
+              + " multiple tokens are valid for a registration. Use an empty string to clear all"
+              + " present default tokens.")
+  List<String> defaultTokens;
+
   /** Returns the existing registry (for update) or null (for creates). */
   @Nullable
   abstract Registry getOldRegistry(String tld);
@@ -373,6 +385,13 @@ abstract class CreateOrUpdateTldCommand extends MutatingCommand {
 
       builder.setAllowedFullyQualifiedHostNames(getAllowedNameservers(oldRegistry));
 
+      if (!isNullOrEmpty(defaultTokens)) {
+        if (defaultTokens.equals(ImmutableList.of(""))) {
+          builder.setDefaultPromoTokens(ImmutableList.of());
+        } else {
+          builder.setDefaultPromoTokens(getTokenKeys(defaultTokens, null));
+        }
+      }
       // Update the Registry object.
       setCommandSpecificProperties(builder);
       stageEntityChange(oldRegistry, builder.build());

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

@@ -24,6 +24,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.persistence.VKey;
@@ -48,11 +49,11 @@ final class DeleteAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
   private static final int BATCH_SIZE = 20;
   private static final Joiner JOINER = Joiner.on(", ");
 
-  private ImmutableSet<VKey<AllocationToken>> tokensToDelete;
+  private ImmutableList<VKey<AllocationToken>> tokensToDelete;
 
   @Override
   public void init() {
-    tokensToDelete = getTokenKeys();
+    tokensToDelete = getTokenKeys(tokens, prefix);
   }
 
   @Override

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

@@ -129,7 +129,7 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     tokensToSave =
         tm().transact(
                 () ->
-                    tm().loadByKeys(getTokenKeys()).values().stream()
+                    tm().loadByKeys(getTokenKeys(tokens, prefix)).values().stream()
                         .collect(toImmutableMap(Function.identity(), this::updateToken))
                         .entrySet()
                         .stream()

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

@@ -16,14 +16,15 @@ package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.beust.jcommander.Parameter;
-import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableList;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.persistence.VKey;
 import java.util.List;
+import javax.annotation.Nullable;
 
 /** Shared base class for commands to update or delete allocation tokens. */
 abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand {
@@ -47,19 +48,20 @@ abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand {
       description = "Do not actually update or delete the tokens; defaults to false")
   protected boolean dryRun;
 
-  protected ImmutableSet<VKey<AllocationToken>> getTokenKeys() {
+  public static ImmutableList<VKey<AllocationToken>> getTokenKeys(
+      @Nullable List<String> tokens, @Nullable String prefix) {
     checkArgument(
         tokens == null ^ prefix == null,
         "Must provide one of --tokens or --prefix, not both / neither");
     if (tokens != null) {
-      ImmutableSet<VKey<AllocationToken>> keys =
+      ImmutableList<VKey<AllocationToken>> keys =
           tokens.stream()
               .map(token -> VKey.create(AllocationToken.class, token))
-              .collect(toImmutableSet());
-      ImmutableSet<VKey<AllocationToken>> nonexistentKeys =
+              .collect(toImmutableList());
+      ImmutableList<VKey<AllocationToken>> nonexistentKeys =
           tm().transact(
-                  () -> keys.stream().filter(key -> !tm().exists(key)).collect(toImmutableSet()));
-      checkState(nonexistentKeys.isEmpty(), "Tokens with keys %s did not exist.", nonexistentKeys);
+                  () -> keys.stream().filter(key -> !tm().exists(key)).collect(toImmutableList()));
+      checkState(nonexistentKeys.isEmpty(), "Tokens with keys %s did not exist", nonexistentKeys);
       return keys;
     } else {
       checkArgument(!prefix.isEmpty(), "Provided prefix should not be blank");
@@ -68,7 +70,7 @@ abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand {
                   tm().loadAllOf(AllocationToken.class).stream()
                       .filter(token -> token.getToken().startsWith(prefix))
                       .map(AllocationToken::createVKey)
-                      .collect(toImmutableSet()));
+                      .collect(toImmutableList()));
     }
   }
 }