From 6a272bc8c6f6b79ed16c608466f45e973edad90d Mon Sep 17 00:00:00 2001
From: gbrodman <gbrodman@google.com>
Date: Wed, 29 May 2019 08:13:01 -0700
Subject: [PATCH] Modify DeleteAllocationTokensCommand to have the same input
 structure as UpdateATC

It's dangerous to have a blank prefix delete all tokens and this allows for some code unification.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=250493453
---
 .../tools/DeleteAllocationTokensCommand.java  | 30 ++-------
 .../tools/UpdateAllocationTokensCommand.java  | 63 ++++--------------
 ...UpdateOrDeleteAllocationTokensCommand.java | 65 +++++++++++++++++++
 .../DeleteAllocationTokensCommandTest.java    | 44 +++++++++----
 .../UpdateAllocationTokensCommandTest.java    | 11 +++-
 5 files changed, 124 insertions(+), 89 deletions(-)
 create mode 100644 java/google/registry/tools/UpdateOrDeleteAllocationTokensCommand.java

diff --git a/java/google/registry/tools/DeleteAllocationTokensCommand.java b/java/google/registry/tools/DeleteAllocationTokensCommand.java
index 88719a5dc..6ee184d12 100644
--- a/java/google/registry/tools/DeleteAllocationTokensCommand.java
+++ b/java/google/registry/tools/DeleteAllocationTokensCommand.java
@@ -25,37 +25,24 @@ import com.beust.jcommander.Parameters;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.cmd.Query;
 import google.registry.model.domain.token.AllocationToken;
 import java.util.List;
 
 /**
  * Command to delete unused {@link AllocationToken}s.
  *
- * <p>Allocation tokens that have been redeemed cannot be deleted. To delete a single allocation
- * token, specify the entire token as the prefix.
+ * <p>Note that all multi-use tokens and redeemed single-use tokens cannot be deleted.
  */
 @Parameters(
     separators = " =",
-    commandDescription = "Deletes the unused AllocationTokens with a given prefix.")
-final class DeleteAllocationTokensCommand extends ConfirmingCommand
-    implements CommandWithRemoteApi {
-
-  @Parameter(
-      names = {"-p", "--prefix"},
-      description = "Allocation token prefix; if blank, deletes all unused tokens",
-      required = true)
-  private String prefix;
+    commandDescription =
+        "Deletes the unused AllocationTokens with a given prefix (or specified tokens).")
+final class DeleteAllocationTokensCommand extends UpdateOrDeleteAllocationTokensCommand {
 
   @Parameter(
       names = {"--with_domains"},
       description = "Allow deletion of allocation tokens with specified domains; defaults to false")
-  boolean withDomains;
-
-  @Parameter(
-      names = {"--dry_run"},
-      description = "Do not actually delete the tokens; defaults to false")
-  boolean dryRun;
+  private boolean withDomains;
 
   private static final int BATCH_SIZE = 20;
   private static final Joiner JOINER = Joiner.on(", ");
@@ -64,12 +51,7 @@ final class DeleteAllocationTokensCommand extends ConfirmingCommand
 
   @Override
   public void init() {
-    Query<AllocationToken> query =
-        ofy().load().type(AllocationToken.class).filter("redemptionHistoryEntry", null);
-    tokensToDelete =
-        query.keys().list().stream()
-            .filter(key -> key.getName().startsWith(prefix))
-            .collect(toImmutableSet());
+    tokensToDelete = getTokenKeys();
   }
 
   @Override
diff --git a/java/google/registry/tools/UpdateAllocationTokensCommand.java b/java/google/registry/tools/UpdateAllocationTokensCommand.java
index 8455568ee..171b9e760 100644
--- a/java/google/registry/tools/UpdateAllocationTokensCommand.java
+++ b/java/google/registry/tools/UpdateAllocationTokensCommand.java
@@ -14,8 +14,6 @@
 
 package google.registry.tools;
 
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Iterables.partition;
@@ -28,7 +26,6 @@ import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
-import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.tools.params.TransitionListParameter.TokenStatusTransitions;
@@ -42,49 +39,30 @@ import org.joda.time.DateTime;
 @Parameters(
     separators = " =",
     commandDescription =
-        "Updates AllocationTokens with the given prefix, modifying some subset of their allowed"
-            + " client IDs, allowed TLDs, discount fraction, or status transitions")
-public final class UpdateAllocationTokensCommand extends ConfirmingCommand
-    implements CommandWithRemoteApi {
-
-  @Parameter(
-      names = {"-p", "--prefix"},
-      description =
-          "Update all allocation tokens with this prefix; otherwise use '--tokens' to specify"
-              + " exact tokens(s) to update")
-  private String prefix;
-
-  @Parameter(
-      names = {"--tokens"},
-      description =
-          "Comma-separated list of tokens to update; otherwise use '--prefix' to update all tokens"
-              + " with a given prefix")
-  private List<String> tokens;
-
-  @Parameter(
-      names = {"--dry_run"},
-      description = "Do not actually update the tokens; defaults to false")
-  private boolean dryRun;
+        "Updates AllocationTokens with the given prefix (or specified tokens), modifying some "
+            + "subset of their allowed client IDs, allowed TLDs, discount fraction, or status "
+            + "transitions")
+final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokensCommand {
 
   @Parameter(
       names = {"--allowed_client_ids"},
       description =
-          "Comma-separated list of allowed client IDs. Use the empty string to clear the"
-              + " existing list.")
+          "Comma-separated list of allowed client IDs. Use the empty string to clear the "
+              + "existing list.")
   private List<String> allowedClientIds;
 
   @Parameter(
       names = {"--allowed_tlds"},
       description =
-          "Comma-separated list of allowed TLDs. Use the empty string to clear the"
-              + " existing list.")
+          "Comma-separated list of allowed TLDs. Use the empty string to clear the "
+              + "existing list.")
   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.")
+          "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(
@@ -92,8 +70,8 @@ public final class UpdateAllocationTokensCommand extends ConfirmingCommand
       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.")
+          "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;
 
   private static final int BATCH_SIZE = 20;
@@ -112,9 +90,8 @@ public final class UpdateAllocationTokensCommand extends ConfirmingCommand
       allowedTlds = ImmutableList.of();
     }
 
-    ImmutableList<Key<AllocationToken>> keysToUpdate = getKeysToUpdate();
     tokensToSave =
-        ofy().load().keys(keysToUpdate).values().stream()
+        ofy().load().keys(getTokenKeys()).values().stream()
             .collect(toImmutableMap(Function.identity(), this::updateToken))
             .entrySet()
             .stream()
@@ -137,20 +114,6 @@ public final class UpdateAllocationTokensCommand extends ConfirmingCommand
     return String.format("Updated %d tokens in total.", numUpdated);
   }
 
-  private ImmutableList<Key<AllocationToken>> getKeysToUpdate() {
-    checkArgument(
-        tokens == null ^ prefix == null, "Must provide one of --tokens or --prefix, not both");
-    if (tokens != null) {
-      return tokens.stream()
-          .map(token -> Key.create(AllocationToken.class, token))
-          .collect(toImmutableList());
-    } else {
-      return ofy().load().type(AllocationToken.class).keys().list().stream()
-          .filter(key -> key.getName().startsWith(prefix))
-          .collect(toImmutableList());
-    }
-  }
-
   private AllocationToken updateToken(AllocationToken original) {
     AllocationToken.Builder builder = original.asBuilder();
     Optional.ofNullable(allowedClientIds)
diff --git a/java/google/registry/tools/UpdateOrDeleteAllocationTokensCommand.java b/java/google/registry/tools/UpdateOrDeleteAllocationTokensCommand.java
new file mode 100644
index 000000000..bd400f447
--- /dev/null
+++ b/java/google/registry/tools/UpdateOrDeleteAllocationTokensCommand.java
@@ -0,0 +1,65 @@
+// 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.tools;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+
+import com.beust.jcommander.Parameter;
+import com.google.common.collect.ImmutableSet;
+import com.googlecode.objectify.Key;
+import google.registry.model.domain.token.AllocationToken;
+import java.util.List;
+
+/** Shared base class for commands to update or delete allocation tokens. */
+abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
+
+  @Parameter(
+      names = {"-p", "--prefix"},
+      description =
+          "Act on all allocation tokens with this prefix, otherwise use '--tokens' to specify "
+              + "exact tokens(s) to act on")
+  protected String prefix;
+
+  @Parameter(
+      names = {"--tokens"},
+      description =
+          "Comma-separated list of tokens to act on; otherwise use '--prefix' to act on all tokens "
+              + "with a given prefix")
+  protected List<String> tokens;
+
+  @Parameter(
+      names = {"--dry_run"},
+      description = "Do not actually update or delete the tokens; defaults to false")
+  protected boolean dryRun;
+
+  protected ImmutableSet<Key<AllocationToken>> getTokenKeys() {
+    checkArgument(
+        tokens == null ^ prefix == null,
+        "Must provide one of --tokens or --prefix, not both / neither");
+    if (tokens != null) {
+      return tokens.stream()
+          .map(token -> Key.create(AllocationToken.class, token))
+          .collect(toImmutableSet());
+    } else {
+      checkArgument(!prefix.isEmpty(), "Provided prefix should not be blank");
+      return ofy().load().type(AllocationToken.class).keys().list().stream()
+          .filter(key -> key.getName().startsWith(prefix))
+          .collect(toImmutableSet());
+    }
+  }
+}
diff --git a/javatests/google/registry/tools/DeleteAllocationTokensCommandTest.java b/javatests/google/registry/tools/DeleteAllocationTokensCommandTest.java
index 94fe5041e..e06b1f4a6 100644
--- a/javatests/google/registry/tools/DeleteAllocationTokensCommandTest.java
+++ b/javatests/google/registry/tools/DeleteAllocationTokensCommandTest.java
@@ -20,7 +20,6 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.JUnitBackports.assertThrows;
 
-import com.beust.jcommander.ParameterException;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenType;
@@ -51,14 +50,6 @@ public class DeleteAllocationTokensCommandTest
     othrNot = persistToken("asdgfho7HASDS", null, false);
   }
 
-  @Test
-  public void test_deleteAllUnredeemedTokens_whenEmptyPrefixSpecified() throws Exception {
-    runCommandForced("--prefix", "");
-    assertThat(reloadTokens(preNot1, preNot2, othrNot)).isEmpty();
-    assertThat(reloadTokens(preRed1, preRed2, othrRed)).containsExactly(preRed1, preRed2, othrRed);
-    assertInStdout("Deleted tokens: asdgfho7HASDS, prefix2978204, prefix8ZZZhs8");
-  }
-
   @Test
   public void test_deleteOnlyUnredeemedTokensWithPrefix() throws Exception {
     runCommandForced("--prefix", "prefix");
@@ -75,6 +66,14 @@ public class DeleteAllocationTokensCommandTest
         .containsExactly(preRed1, preRed2, preNot1, preNot2, othrRed);
   }
 
+  @Test
+  public void test_deleteParticularTokens() throws Exception {
+    runCommandForced("--tokens", "prefix2978204,asdgfho7HASDS");
+    assertThat(reloadTokens(preNot1, othrNot)).isEmpty();
+    assertThat(reloadTokens(preRed1, preRed2, preNot2, othrRed))
+        .containsExactly(preRed1, preRed2, preNot2, othrRed);
+  }
+
   @Test
   public void test_deleteTokensWithNonExistentPrefix_doesNothing() throws Exception {
     runCommandForced("--prefix", "nonexistent");
@@ -84,10 +83,10 @@ public class DeleteAllocationTokensCommandTest
 
   @Test
   public void test_dryRun_deletesNothing() throws Exception {
-    runCommandForced("--prefix", "", "--dry_run");
+    runCommandForced("--prefix", "prefix", "--dry_run");
     assertThat(reloadTokens(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot))
         .containsExactly(preRed1, preRed2, preNot1, preNot2, othrRed, othrNot);
-    assertInStdout("Would delete tokens: asdgfho7HASDS, prefix2978204, prefix8ZZZhs8");
+    assertInStdout("Would delete tokens: prefix2978204, prefix8ZZZhs8");
   }
 
   @Test
@@ -134,10 +133,29 @@ public class DeleteAllocationTokensCommandTest
 
   @Test
   public void test_prefixIsRequired() {
-    ParameterException thrown = assertThrows(ParameterException.class, () -> runCommandForced());
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, this::runCommandForced);
     assertThat(thrown)
         .hasMessageThat()
-        .isEqualTo("The following option is required: -p, --prefix ");
+        .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
+  }
+
+  @Test
+  public void testFailure_bothPrefixAndTokens() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> runCommandForced("--prefix", "somePrefix", "--tokens", "someToken"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
+  }
+
+  @Test
+  public void testFailure_emptyPrefix() {
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, () -> runCommandForced("--prefix", ""));
+    assertThat(thrown).hasMessageThat().isEqualTo("Provided prefix should not be blank");
   }
 
   private static AllocationToken persistToken(
diff --git a/javatests/google/registry/tools/UpdateAllocationTokensCommandTest.java b/javatests/google/registry/tools/UpdateAllocationTokensCommandTest.java
index 447967e19..90c5f8d02 100644
--- a/javatests/google/registry/tools/UpdateAllocationTokensCommandTest.java
+++ b/javatests/google/registry/tools/UpdateAllocationTokensCommandTest.java
@@ -176,7 +176,7 @@ public class UpdateAllocationTokensCommandTest
                 IllegalArgumentException.class,
                 () -> runCommandForced("--prefix", "token", "--tokens", "token")))
         .hasMessageThat()
-        .isEqualTo("Must provide one of --tokens or --prefix, not both");
+        .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
   }
 
   @Test
@@ -185,7 +185,14 @@ public class UpdateAllocationTokensCommandTest
             assertThrows(
                 IllegalArgumentException.class, () -> runCommandForced("--allowed_tlds", "tld")))
         .hasMessageThat()
-        .isEqualTo("Must provide one of --tokens or --prefix, not both");
+        .isEqualTo("Must provide one of --tokens or --prefix, not both / neither");
+  }
+
+  @Test
+  public void testFailure_emptyPrefix() {
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, () -> runCommandForced("--prefix", ""));
+    assertThat(thrown).hasMessageThat().isEqualTo("Provided prefix should not be blank");
   }
 
   private static AllocationToken.Builder builderWithPromo() {