From bb27feebd62080a1632bad45cfe2ba6ca1b5e4ac Mon Sep 17 00:00:00 2001
From: Rachel Guan <rachelguan@google.com>
Date: Thu, 5 May 2022 15:48:45 -0400
Subject: [PATCH] Add renewal logic in allocation token related commands 
 (#1596)

* Add renewal price behavior to allocation token related command

* Add details to renewal price behavior
---
 .../GenerateAllocationTokensCommand.java      | 13 +++
 .../tools/UpdateAllocationTokensCommand.java  | 15 +++
 .../domain/token/AllocationTokenTest.java     |  6 +-
 .../GenerateAllocationTokensCommandTest.java  | 69 ++++++++++++++
 .../UpdateAllocationTokensCommandTest.java    | 95 +++++++++++++++++++
 5 files changed, 195 insertions(+), 3 deletions(-)

diff --git a/core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java b/core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java
index 16137fddc..d1d9cc909 100644
--- a/core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java
+++ b/core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java
@@ -17,6 +17,7 @@ package google.registry.tools;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Sets.difference;
+import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.DEFAULT;
 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.persistence.transaction.TransactionManagerFactory.tm;
@@ -37,6 +38,7 @@ import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Streams;
 import com.google.common.io.Files;
+import google.registry.model.billing.BillingEvent.RenewalPriceBehavior;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.domain.token.AllocationToken.TokenType;
@@ -141,6 +143,16 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
               + " form <time>=<status>[,<time>=<status>]* where each status represents the status.")
   private ImmutableSortedMap<DateTime, TokenStatus> tokenStatusTransitions;
 
+  @Parameter(
+      names = {"--renewal_price_behavior"},
+      description =
+          "The type of renewal price behavior, either DEFAULT (default), NONPREMIUM, or SPECIFIED."
+              + " This indicates how a domain should be charged for renewal. By default, a domain"
+              + " will be renewed at the renewal price from the pricing engine. If the renewal"
+              + " price behavior is set to SPECIFIED, it means that the renewal cost will be the"
+              + " same as the domain's calculated create price.")
+  private RenewalPriceBehavior renewalPriceBehavior = DEFAULT;
+
   @Parameter(
       names = {"--dry_run"},
       description = "Do not actually persist the tokens; defaults to false")
@@ -184,6 +196,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
                     AllocationToken.Builder token =
                         new AllocationToken.Builder()
                             .setToken(t)
+                            .setRenewalPriceBehavior(renewalPriceBehavior)
                             .setTokenType(tokenType == null ? SINGLE_USE : tokenType)
                             .setAllowedRegistrarIds(
                                 ImmutableSet.copyOf(nullToEmpty(allowedClientIds)))
diff --git a/core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java b/core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java
index 5e2699a60..3e0e820dd 100644
--- a/core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java
+++ b/core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java
@@ -23,10 +23,12 @@ import static google.registry.persistence.transaction.TransactionManagerUtil.tra
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import com.beust.jcommander.internal.Nullable;
 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 google.registry.model.billing.BillingEvent.RenewalPriceBehavior;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.tools.params.TransitionListParameter.TokenStatusTransitions;
@@ -88,6 +90,17 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
               + "form <time>=<status>[,<time>=<status>]* where each status represents the status.")
   private ImmutableSortedMap<DateTime, TokenStatus> tokenStatusTransitions;
 
+  @Parameter(
+      names = {"--renewal_price_behavior"},
+      description =
+          "The type of renewal price behavior, either DEFAULT (default), NONPREMIUM, or SPECIFIED."
+              + " This indicates how a domain should be charged for renewal. By default, a domain"
+              + " will be renewed at the renewal price from the pricing engine. If the renewal"
+              + " price behavior is set to SPECIFIED, it means that the renewal cost will be the"
+              + " same as the domain's calculated create price.")
+  @Nullable
+  private RenewalPriceBehavior renewalPriceBehavior;
+
   private static final int BATCH_SIZE = 20;
   private static final Joiner JOINER = Joiner.on(", ");
 
@@ -142,6 +155,8 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
     Optional.ofNullable(discountPremiums).ifPresent(builder::setDiscountPremiums);
     Optional.ofNullable(discountYears).ifPresent(builder::setDiscountYears);
     Optional.ofNullable(tokenStatusTransitions).ifPresent(builder::setTokenStatusTransitions);
+    Optional.ofNullable(renewalPriceBehavior)
+        .ifPresent(behavior -> builder.setRenewalPriceBehavior(renewalPriceBehavior));
     return builder.build();
   }
 
diff --git a/core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java b/core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java
index db8c579e9..6b4738d7b 100644
--- a/core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java
+++ b/core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java
@@ -162,7 +162,7 @@ public class AllocationTokenTest extends EntityTestCase {
   }
 
   @TestOfyAndSql
-  void testgetRenewalBehavior_returnsDefaultRenewBehavior() {
+  void testGetRenewalBehavior_returnsDefaultRenewBehavior() {
     assertThat(
             persistResource(
                     new AllocationToken.Builder()
@@ -174,7 +174,7 @@ public class AllocationTokenTest extends EntityTestCase {
   }
 
   @TestOfyAndSql
-  void testsetRenewalBehavior_assertsRenewalBehaviorIsNotDefault() {
+  void testSetRenewalBehavior_assertsRenewalBehaviorIsNotDefault() {
     assertThat(
             persistResource(
                     new AllocationToken.Builder()
@@ -187,7 +187,7 @@ public class AllocationTokenTest extends EntityTestCase {
   }
 
   @TestOfyAndSql
-  void testsetRenewalBehavior_assertRenewalBehaviorIsModified() {
+  void testSetRenewalBehavior_assertRenewalBehaviorIsModified() {
     AllocationToken token =
         persistResource(
             new AllocationToken.Builder()
diff --git a/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java b/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java
index 4b7cd98f2..215dae38a 100644
--- a/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java
+++ b/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java
@@ -15,6 +15,8 @@
 package google.registry.tools;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.NONPREMIUM;
+import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.SPECIFIED;
 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.testing.DatabaseHelper.assertAllocationTokens;
@@ -193,6 +195,73 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
     assertInStdout("foobar", "foobaz");
   }
 
+  @TestOfyAndSql
+  void testSuccess_renewalPriceBehaviorIsDefault() throws Exception {
+    runCommand("--tokens", "foobar,foobaz", "--renewal_price_behavior", "DEFAULT");
+    assertAllocationTokens(createToken("foobar", null, null), createToken("foobaz", null, null));
+    assertInStdout("foobar", "foobaz");
+  }
+
+  @TestOfyAndSql
+  void testSuccess_renewalPriceBehaviorIsSetToDefaultByDefault() throws Exception {
+    runCommand("--tokens", "foobar,foobaz");
+    assertAllocationTokens(createToken("foobar", null, null), createToken("foobaz", null, null));
+    assertInStdout("foobar", "foobaz");
+  }
+
+  @TestOfyAndSql
+  void testSuccess_renewalPriceBehaviorIsNonPremium() throws Exception {
+    runCommand("--tokens", "foobar,foobaz", "--renewal_price_behavior", "NONPREMIUM");
+    assertAllocationTokens(
+        createToken("foobar", null, null).asBuilder().setRenewalPriceBehavior(NONPREMIUM).build(),
+        createToken("foobaz", null, null).asBuilder().setRenewalPriceBehavior(NONPREMIUM).build());
+    assertInStdout("foobar", "foobaz");
+  }
+
+  @TestOfyAndSql
+  void testSuccess_renewalPriceBehaviorIsSpecified() throws Exception {
+    runCommand("--tokens", "foobar,foobaz", "--renewal_price_behavior", "SPECIFIED");
+    assertAllocationTokens(
+        createToken("foobar", null, null).asBuilder().setRenewalPriceBehavior(SPECIFIED).build(),
+        createToken("foobaz", null, null).asBuilder().setRenewalPriceBehavior(SPECIFIED).build());
+    assertInStdout("foobar", "foobaz");
+  }
+
+  @TestOfyAndSql
+  void testSuccess_renewalPriceBehaviorIsSpecifiedButMixedCase() throws Exception {
+    runCommand("--tokens", "foobar,foobaz", "--renewal_price_behavior", "speCIFied");
+    assertAllocationTokens(
+        createToken("foobar", null, null).asBuilder().setRenewalPriceBehavior(SPECIFIED).build(),
+        createToken("foobaz", null, null).asBuilder().setRenewalPriceBehavior(SPECIFIED).build());
+    assertInStdout("foobar", "foobaz");
+  }
+
+  @TestOfyAndSql
+  void testFailure_renewalPriceBehaviorIsInvalid() {
+    ParameterException thrown =
+        assertThrows(
+            ParameterException.class,
+            () -> runCommand("--tokens", "foobar,foobaz", "--renewal_price_behavior", "SPEXIFIED"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "Invalid value for --renewal_price_behavior parameter. Allowed values:[DEFAULT,"
+                + " NONPREMIUM, SPECIFIED]");
+  }
+
+  @TestOfyAndSql
+  void testFailure_renewalPriceBehaviorIsEmptyString() {
+    ParameterException thrown =
+        assertThrows(
+            ParameterException.class,
+            () -> runCommand("--tokens", "foobar,foobaz", "--renewal_price_behavior", ""));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "Invalid value for --renewal_price_behavior parameter. Allowed values:[DEFAULT,"
+                + " NONPREMIUM, SPECIFIED]");
+  }
+
   @TestOfyAndSql
   void testSuccess_specifyManyTokens() throws Exception {
     command.stringGenerator =
diff --git a/core/src/test/java/google/registry/tools/UpdateAllocationTokensCommandTest.java b/core/src/test/java/google/registry/tools/UpdateAllocationTokensCommandTest.java
index 7615b59d2..b8db4dfb2 100644
--- a/core/src/test/java/google/registry/tools/UpdateAllocationTokensCommandTest.java
+++ b/core/src/test/java/google/registry/tools/UpdateAllocationTokensCommandTest.java
@@ -15,6 +15,9 @@
 package google.registry.tools;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.DEFAULT;
+import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.NONPREMIUM;
+import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.SPECIFIED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.CANCELLED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.ENDED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.NOT_STARTED;
@@ -26,6 +29,7 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static org.joda.time.DateTimeZone.UTC;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.beust.jcommander.ParameterException;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import google.registry.model.domain.token.AllocationToken;
@@ -34,6 +38,7 @@ import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
 
+/** Unit tests for {@link UpdateAllocationTokensCommand}. */
 @DualDatabaseTest
 class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocationTokensCommand> {
 
@@ -97,6 +102,96 @@ class UpdateAllocationTokensCommandTest extends CommandTestCase<UpdateAllocation
     assertThat(reloadResource(token).getDiscountYears()).isEqualTo(4);
   }
 
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setToSpecified() throws Exception {
+    AllocationToken token = persistResource(builderWithPromo().setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--renewal_price_behavior", "SPECIFIED");
+    assertThat(reloadResource(token).getRenewalPriceBehavior()).isEqualTo(SPECIFIED);
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setToDefault() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setRenewalPriceBehavior(SPECIFIED).setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--renewal_price_behavior", "default");
+    assertThat(reloadResource(token).getRenewalPriceBehavior()).isEqualTo(DEFAULT);
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setToNonPremium() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setRenewalPriceBehavior(SPECIFIED).setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--renewal_price_behavior", "NONpremium");
+    assertThat(reloadResource(token).getRenewalPriceBehavior()).isEqualTo(NONPREMIUM);
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setFromDefaultToDefault() throws Exception {
+    AllocationToken token = persistResource(builderWithPromo().setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--renewal_price_behavior", "defauLT");
+    assertThat(reloadResource(token).getRenewalPriceBehavior()).isEqualTo(DEFAULT);
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setFromSpecifiedToSpecified() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setRenewalPriceBehavior(SPECIFIED).setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--renewal_price_behavior", "SPecified");
+    assertThat(reloadResource(token).getRenewalPriceBehavior()).isEqualTo(SPECIFIED);
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setFromNonPremiumToDefault() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo()
+                .setRenewalPriceBehavior(NONPREMIUM)
+                .setDiscountFraction(0.5)
+                .build());
+    runCommandForced("--prefix", "token", "--renewal_price_behavior", "defauLT");
+    assertThat(reloadResource(token).getRenewalPriceBehavior()).isEqualTo(DEFAULT);
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setToMixedCaseDefault() throws Exception {
+    AllocationToken token =
+        persistResource(
+            builderWithPromo().setRenewalPriceBehavior(SPECIFIED).setDiscountFraction(0.5).build());
+    runCommandForced("--prefix", "token", "--renewal_price_behavior", "deFauLt");
+    assertThat(reloadResource(token).getRenewalPriceBehavior()).isEqualTo(DEFAULT);
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setToInvalidBehavior_throwsException() {
+    ParameterException thrown =
+        assertThrows(
+            ParameterException.class,
+            () -> runCommandForced("--prefix", "token", "--renewal_price_behavior", "premium"));
+    persistResource(builderWithPromo().setDiscountFraction(0.5).build());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "Invalid value for --renewal_price_behavior parameter. Allowed values:[DEFAULT,"
+                + " NONPREMIUM, SPECIFIED]");
+  }
+
+  @TestOfyAndSql
+  void testUpdateRenewalPriceBehavior_setToEmptyString_throwsException() {
+    ParameterException thrown =
+        assertThrows(
+            ParameterException.class,
+            () -> runCommandForced("--prefix", "token", "--renewal_price_behavior", ""));
+    persistResource(builderWithPromo().setDiscountFraction(0.5).build());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo(
+            "Invalid value for --renewal_price_behavior parameter. Allowed values:[DEFAULT,"
+                + " NONPREMIUM, SPECIFIED]");
+  }
+
   @TestOfyAndSql
   void testUpdateStatusTransitions() throws Exception {
     DateTime now = DateTime.now(UTC);