Add currentPackageToken on create flow (#1751)

* Add currentPackageToken on create flow

* Change to Truth8 assertion

* Add check for specified renewal behavior

core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java

@@ -390,6 +390,11 @@ public final class DomainCreateFlow implements TransactionalFlow {
             .addGracePeriod(
                 GracePeriod.forBillingEvent(GracePeriodStatus.ADD, repoId, createBillingEvent))
             .build();
+    if (allocationToken.isPresent()
+        && allocationToken.get().getTokenType().equals(TokenType.PACKAGE)) {
+      domain =
+          domain.asBuilder().setCurrentPackageToken(allocationToken.get().createVKey()).build();
+    }
     DomainHistory domainHistory =
         buildDomainHistory(domain, registry, now, period, registry.getAddGracePeriodLength());
     if (reservationTypes.contains(NAME_COLLISION)) {

core/src/main/java/google/registry/model/domain/token/AllocationToken.java

@@ -105,6 +105,7 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
 
   /** Single-use tokens are invalid after use. Infinite-use tokens, predictably, are not. */
   public enum TokenType {
+    PACKAGE,
     SINGLE_USE,
     UNLIMITED_USE
   }
@@ -274,6 +275,10 @@ public class AllocationToken extends BackupGroupRoot implements Buildable {
     public AllocationToken build() {
       checkArgumentNotNull(getInstance().tokenType, "Token type must be specified");
       checkArgument(!Strings.isNullOrEmpty(getInstance().token), "Token must not be null or empty");
+      checkArgument(
+          !getInstance().tokenType.equals(TokenType.PACKAGE)
+              || getInstance().renewalPriceBehavior.equals(RenewalPriceBehavior.SPECIFIED),
+          "Package tokens must have renewalPriceBehavior set to SPECIFIED");
       checkArgument(
           getInstance().domainName == null || TokenType.SINGLE_USE.equals(getInstance().tokenType),
           "Domain name can only be specified for SINGLE_USE tokens");

core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -25,6 +25,7 @@ import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.DE
 import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.NONPREMIUM;
 import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.SPECIFIED;
 import static google.registry.model.domain.fee.Fee.FEE_EXTENSION_URIS;
+import static google.registry.model.domain.token.AllocationToken.TokenType.PACKAGE;
 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.eppcommon.StatusValue.PENDING_DELETE;
@@ -70,6 +71,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
+import com.google.common.truth.Truth8;
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig;
 import google.registry.flows.EppException;
@@ -3126,4 +3128,25 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         ImmutableMap.of("DOMAIN", "example-one.tld", "YEARS", "2"));
     assertThrows(AllocationTokenNotValidForDomainException.class, this::runFlow);
   }
+
+  @Test
+  void testSuccess_packageToken_addsTokenToDomain() throws Exception {
+    AllocationToken token =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("abc123")
+                .setTokenType(PACKAGE)
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .setRenewalPriceBehavior(SPECIFIED)
+                .build());
+    persistContactsAndHosts();
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
+    runFlowAssertResponse(
+        loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.tld")));
+    Domain domain = reloadResourceByForeignKey();
+    assertThat(domain.getCurrentPackageToken()).isPresent();
+    Truth8.assertThat(domain.getCurrentPackageToken()).hasValue(token.createVKey());
+  }
 }

core/src/test/java/google/registry/model/domain/token/AllocationTokenTest.java

@@ -213,6 +213,20 @@ public class AllocationTokenTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().isEqualTo("Token type can only be set once");
   }
 
+  @Test
+  void testFail_packageTokenNotSpecifiedRenewalBehavior() {
+    AllocationToken.Builder builder =
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(TokenType.PACKAGE)
+            .setRenewalPriceBehavior(RenewalPriceBehavior.DEFAULT);
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, () -> builder.build());
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Package tokens must have renewalPriceBehavior set to SPECIFIED");
+  }
+
   @Test
   void testBuild_DomainNameWithLessThanTwoParts() {
     IllegalArgumentException thrown =

core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java

@@ -402,7 +402,8 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase<GenerateAlloca
             () -> runCommand("--number", "999", "--type", "INVALID_TYPE"));
     assertThat(thrown)
         .hasMessageThat()
-        .isEqualTo("Invalid value for -t parameter. Allowed values:[SINGLE_USE, UNLIMITED_USE]");
+        .isEqualTo(
+            "Invalid value for -t parameter. Allowed values:[PACKAGE, SINGLE_USE, UNLIMITED_USE]");
   }
 
   @Test

core/src/test/resources/google/registry/model/schema.txt

@@ -258,6 +258,7 @@ enum google.registry.model.domain.token.AllocationToken$TokenStatus {
   VALID;
 }
 enum google.registry.model.domain.token.AllocationToken$TokenType {
+  PACKAGE;
   SINGLE_USE;
   UNLIMITED_USE;
 }