Store a reference to an allocation token in the OneTime billing event

We will need to be able to find all redemptions associated with a particular token and this will allow us to do that.

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

java/google/registry/flows/domain/DomainCreateFlow.java

@@ -295,6 +295,7 @@ public class DomainCreateFlow implements TransactionalFlow {
             years,
             feesAndCredits,
             historyEntry,
+            allocationToken,
             now);
     // Create a new autorenew billing event and poll message starting at the expiration time.
     BillingEvent.Recurring autorenewBillingEvent =
@@ -477,6 +478,7 @@ public class DomainCreateFlow implements TransactionalFlow {
       int years,
       FeesAndCredits feesAndCredits,
       HistoryEntry historyEntry,
+      Optional<AllocationToken> allocationToken,
       DateTime now) {
     ImmutableSet.Builder<Flag> flagsBuilder = new ImmutableSet.Builder<>();
     // Sunrise and anchor tenancy are orthogonal tags and thus both can be present together.
@@ -497,6 +499,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         .setPeriodYears(years)
         .setCost(feesAndCredits.getCreateCost())
         .setEventTime(now)
+        .setAllocationToken(allocationToken.map(Key::create).orElse(null))
         .setBillingTime(
             now.plus(
                 isAnchorTenant

java/google/registry/model/billing/BillingEvent.java

@@ -40,6 +40,7 @@ import google.registry.model.annotations.ReportedOn;
 import google.registry.model.common.TimeOfYear;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
 import java.util.Objects;
@@ -244,6 +245,11 @@ public abstract class BillingEvent extends ImmutableObject
      */
     Key<? extends BillingEvent> cancellationMatchingBillingEvent;
 
+    /**
+     * The {@link AllocationToken} used in the creation of this event, or null if one was not used.
+     */
+    @Index @Nullable Key<AllocationToken> allocationToken;
+
     public Money getCost() {
       return cost;
     }
@@ -264,6 +270,10 @@ public abstract class BillingEvent extends ImmutableObject
       return cancellationMatchingBillingEvent;
     }
 
+    public Optional<Key<AllocationToken>> getAllocationToken() {
+      return Optional.ofNullable(allocationToken);
+    }
+
     @Override
     public Builder asBuilder() {
       return new Builder(clone(this));
@@ -306,6 +316,11 @@ public abstract class BillingEvent extends ImmutableObject
         return this;
       }
 
+      public Builder setAllocationToken(@Nullable Key<AllocationToken> allocationToken) {
+        getInstance().allocationToken = allocationToken;
+        return this;
+      }
+
       @Override
       public OneTime build() {
         OneTime instance = getInstance();

javatests/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.flows.FlowTestCase.UserPrivileges.SUPERUSER;
 import static google.registry.model.billing.BillingEvent.Flag.ANCHOR_TENANT;
+import static google.registry.model.billing.BillingEvent.Flag.RESERVED;
 import static google.registry.model.billing.BillingEvent.Flag.SUNRISE;
 import static google.registry.model.domain.fee.Fee.FEE_EXTENSION_URIS;
 import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
@@ -151,6 +152,7 @@ import google.registry.model.reporting.HistoryEntry;
 import google.registry.monitoring.whitebox.EppMetric;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import java.util.Map;
+import javax.annotation.Nullable;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
@@ -162,6 +164,8 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
 
   private static final String CLAIMS_KEY = "2013041500/2/6/9/rJ1NrDO92vDsAzf7EQzgjX4R0000000001";
 
+  private AllocationToken allocationToken;
+
   public DomainCreateFlowTest() {
     setEppInput("domain_create.xml", ImmutableMap.of("DOMAIN", "example.tld"));
     clock.setTo(DateTime.parse("1999-04-03T22:00:00.0Z").minus(Duration.millis(1)));
@@ -170,12 +174,13 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
   @Before
   public void initCreateTest() {
     createTld("tld");
-    persistResource(
+    allocationToken =
-        new AllocationToken.Builder()
+        persistResource(
-            .setToken("abcDEF23456")
+            new AllocationToken.Builder()
-            .setTokenType(SINGLE_USE)
+                .setToken("abcDEF23456")
-            .setDomainName("anchor.tld")
+                .setTokenType(SINGLE_USE)
-            .build());
+                .setDomainName("anchor.tld")
+                .build());
     persistResource(
         Registry.get("tld")
             .asBuilder()
@@ -212,6 +217,14 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
 
   private void assertSuccessfulCreate(
       String domainTld, ImmutableSet<BillingEvent.Flag> expectedBillingFlags) throws Exception {
+    assertSuccessfulCreate(domainTld, expectedBillingFlags, null);
+  }
+
+  private void assertSuccessfulCreate(
+      String domainTld,
+      ImmutableSet<BillingEvent.Flag> expectedBillingFlags,
+      @Nullable AllocationToken allocationToken)
+      throws Exception {
     DomainBase domain = reloadResourceByForeignKey();
 
     boolean isAnchorTenant = expectedBillingFlags.contains(ANCHOR_TENANT);
@@ -252,6 +265,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
             .setBillingTime(billingTime)
             .setFlags(expectedBillingFlags)
             .setParent(historyEntry)
+            .setAllocationToken(allocationToken == null ? null : Key.create(allocationToken))
             .build();
 
     BillingEvent.Recurring renewBillingEvent =
@@ -440,7 +454,8 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
         persistResource(
             new AllocationToken.Builder().setToken("abc123").setTokenType(SINGLE_USE).build());
     clock.advanceOneMilli();
-    doSuccessfulTest();
+    runFlow();
+    assertSuccessfulCreate("tld", ImmutableSet.of(), token);
     HistoryEntry historyEntry =
         ofy().load().type(HistoryEntry.class).ancestor(reloadResourceByForeignKey()).first().now();
     assertThat(ofy().load().entity(token).now().getRedemptionHistoryEntry())
@@ -451,10 +466,12 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
   public void testSuccess_validAllocationToken_multiUse() throws Exception {
     setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "example.tld"));
     persistContactsAndHosts();
-    persistResource(
+    allocationToken =
-        new AllocationToken.Builder().setTokenType(UNLIMITED_USE).setToken("abc123").build());
+        persistResource(
+            new AllocationToken.Builder().setTokenType(UNLIMITED_USE).setToken("abc123").build());
     clock.advanceOneMilli();
-    doSuccessfulTest();
+    runFlow();
+    assertSuccessfulCreate("tld", ImmutableSet.of(), allocationToken);
     clock.advanceOneMilli();
     setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "otherexample.tld"));
     runFlowAssertResponse(
@@ -955,7 +972,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
     setEppInput("domain_create_anchor_allocationtoken.xml");
     persistContactsAndHosts();
     runFlowAssertResponse(loadFile("domain_create_anchor_response.xml"));
-    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT));
+    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT), allocationToken);
     assertNoLordn();
     assertAllocationTokenWasRedeemed("abcDEF23456");
   }
@@ -975,7 +992,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
     setEppInput("domain_create_anchor_allocationtoken.xml");
     persistContactsAndHosts();
     runFlowAssertResponse(loadFile("domain_create_anchor_response.xml"));
-    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT));
+    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT), allocationToken);
     assertNoLordn();
     assertAllocationTokenWasRedeemed("abcDEF23456");
   }
@@ -998,7 +1015,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
     clock.setTo(DateTime.parse("2009-08-16T09:00:00.0Z"));
     persistContactsAndHosts();
     runFlowAssertResponse(loadFile("domain_create_response_claims.xml"));
-    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT));
+    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT), allocationToken);
     assertDnsTasksEnqueued("example-one.tld");
     assertClaimsLordn();
     assertAllocationTokenWasRedeemed("abcDEF23456");
@@ -1029,12 +1046,13 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
 
   @Test
   public void testSuccess_anchorTenantInSunrise_withSignedMark() throws Exception {
-    persistResource(
+    allocationToken =
-        new AllocationToken.Builder()
+        persistResource(
-            .setDomainName("test-validate.tld")
+            new AllocationToken.Builder()
-            .setToken("abcDEF23456")
+                .setDomainName("test-validate.tld")
-            .setTokenType(SINGLE_USE)
+                .setToken("abcDEF23456")
-            .build());
+                .setTokenType(SINGLE_USE)
+                .build());
     persistResource(
         Registry.get("tld")
             .asBuilder()
@@ -1049,7 +1067,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
         loadFile(
             "domain_create_response_encoded_signed_mark_name.xml",
             ImmutableMap.of("DOMAIN", "test-validate.tld")));
-    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT, SUNRISE));
+    assertSuccessfulCreate("tld", ImmutableSet.of(ANCHOR_TENANT, SUNRISE), allocationToken);
     assertDnsTasksEnqueued("test-validate.tld");
     assertSunriseLordn("test-validate.tld");
     assertAllocationTokenWasRedeemed("abcDEF23456");
@@ -1057,19 +1075,20 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
 
   @Test
   public void testSuccess_reservedDomain_viaAllocationTokenExtension() throws Exception {
-    persistResource(
+    allocationToken =
-        new AllocationToken.Builder()
+        persistResource(
-            .setToken("abc123")
+            new AllocationToken.Builder()
-            .setTokenType(SINGLE_USE)
+                .setToken("abc123")
-            .setDomainName("resdom.tld")
+                .setTokenType(SINGLE_USE)
-            .build());
+                .setDomainName("resdom.tld")
+                .build());
     // Despite the domain being FULLY_BLOCKED, the non-superuser create succeeds the domain is also
     // RESERVED_FOR_SPECIFIC_USE and the correct allocation token is passed.
     setEppInput("domain_create_allocationtoken.xml", ImmutableMap.of("DOMAIN", "resdom.tld"));
     persistContactsAndHosts();
     runFlowAssertResponse(
         loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "resdom.tld")));
-    assertSuccessfulCreate("tld", ImmutableSet.of(Flag.RESERVED));
+    assertSuccessfulCreate("tld", ImmutableSet.of(RESERVED), allocationToken);
     assertNoLordn();
     assertAllocationTokenWasRedeemed("abc123");
   }
@@ -1088,7 +1107,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
     persistContactsAndHosts();
     runFlowAssertResponse(
         CommitMode.LIVE, SUPERUSER, loadFile("domain_create_reserved_response.xml"));
-    assertSuccessfulCreate("tld", ImmutableSet.of(Flag.RESERVED));
+    assertSuccessfulCreate("tld", ImmutableSet.of(RESERVED));
   }
 
   @Test
@@ -1643,7 +1662,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
         loadFile(
             "domain_create_response_encoded_signed_mark_name.xml",
             ImmutableMap.of("DOMAIN", "test-validate.tld")));
-    assertSuccessfulCreate("tld", ImmutableSet.of(Flag.SUNRISE));
+    assertSuccessfulCreate("tld", ImmutableSet.of(SUNRISE));
     assertSunriseLordn("test-validate.tld");
   }
 
@@ -1661,7 +1680,7 @@ public class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow,
         loadFile(
             "domain_create_response_encoded_signed_mark_name.xml",
             ImmutableMap.of("DOMAIN", "test-validate.tld")));
-    assertSuccessfulCreate("tld", ImmutableSet.of(Flag.SUNRISE));
+    assertSuccessfulCreate("tld", ImmutableSet.of(SUNRISE));
     assertSunriseLordn("test-validate.tld");
   }
 

javatests/google/registry/model/billing/BillingEventTest.java

@@ -15,6 +15,7 @@
 package google.registry.model.billing;
 
 import static com.google.common.truth.Truth.assertThat;
+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.createTld;
 import static google.registry.testing.DatastoreHelper.persistActiveDomain;
@@ -25,6 +26,7 @@ import static org.joda.money.CurrencyUnit.USD;
 import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
 import com.googlecode.objectify.Key;
 import google.registry.model.EntityTestCase;
 import google.registry.model.billing.BillingEvent.Flag;
@@ -32,7 +34,10 @@ import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.model.domain.token.AllocationToken;
+import google.registry.model.domain.token.AllocationToken.TokenStatus;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.util.DateTimeUtils;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.Before;
@@ -67,22 +72,41 @@ public class BillingEventTest extends EntityTestCase {
         .setModificationTime(now.plusDays(1))
         .build());
 
-    oneTime = persistResource(commonInit(
+    AllocationToken allocationToken =
-        new BillingEvent.OneTime.Builder()
+        persistResource(
-            .setParent(historyEntry)
+            new AllocationToken.Builder()
-            .setReason(Reason.CREATE)
+                .setToken("abc123")
-            .setFlags(ImmutableSet.of(BillingEvent.Flag.ANCHOR_TENANT))
+                .setTokenType(UNLIMITED_USE)
-            .setPeriodYears(2)
+                .setDiscountFraction(0.5)
-            .setCost(Money.of(USD, 1))
+                .setTokenStatusTransitions(
-            .setEventTime(now)
+                    ImmutableSortedMap.<DateTime, TokenStatus>naturalOrder()
-            .setBillingTime(now.plusDays(5))));
+                        .put(DateTimeUtils.START_OF_TIME, TokenStatus.NOT_STARTED)
-    recurring = persistResource(commonInit(
+                        .put(DateTime.now(UTC), TokenStatus.VALID)
-        new BillingEvent.Recurring.Builder()
+                        .put(DateTime.now(UTC).plusWeeks(8), TokenStatus.ENDED)
-            .setParent(historyEntry)
+                        .build())
-            .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+                .build());
-            .setReason(Reason.RENEW)
+
-            .setEventTime(now.plusYears(1))
+    oneTime =
-            .setRecurrenceEndTime(END_OF_TIME)));
+        persistResource(
+            commonInit(
+                new BillingEvent.OneTime.Builder()
+                    .setParent(historyEntry)
+                    .setReason(Reason.CREATE)
+                    .setFlags(ImmutableSet.of(BillingEvent.Flag.ANCHOR_TENANT))
+                    .setPeriodYears(2)
+                    .setCost(Money.of(USD, 1))
+                    .setEventTime(now)
+                    .setBillingTime(now.plusDays(5))
+                    .setAllocationToken(Key.create(allocationToken))));
+    recurring =
+        persistResource(
+            commonInit(
+                new BillingEvent.Recurring.Builder()
+                    .setParent(historyEntry)
+                    .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+                    .setReason(Reason.RENEW)
+                    .setEventTime(now.plusYears(1))
+                    .setRecurrenceEndTime(END_OF_TIME)));
     oneTimeSynthetic = persistResource(commonInit(
         new BillingEvent.OneTime.Builder()
             .setParent(historyEntry)
@@ -166,9 +190,20 @@ public class BillingEventTest extends EntityTestCase {
 
   @Test
   public void testIndexing() throws Exception {
-    verifyIndexing(oneTime, "clientId", "eventTime", "billingTime", "syntheticCreationTime");
     verifyIndexing(
-        oneTimeSynthetic, "clientId", "eventTime", "billingTime", "syntheticCreationTime");
+        oneTime,
+        "clientId",
+        "eventTime",
+        "billingTime",
+        "syntheticCreationTime",
+        "allocationToken");
+    verifyIndexing(
+        oneTimeSynthetic,
+        "clientId",
+        "eventTime",
+        "billingTime",
+        "syntheticCreationTime",
+        "allocationToken");
     verifyIndexing(
         recurring, "clientId", "eventTime", "recurrenceEndTime", "recurrenceTimeOfYear.timeString");
     verifyIndexing(cancellationOneTime, "clientId", "eventTime", "billingTime");

javatests/google/registry/model/testdata/schema.txt

@@ -41,6 +41,7 @@ class google.registry.model.billing.BillingEvent$OneTime {
   @Id long id;
   @Parent com.googlecode.objectify.Key<google.registry.model.reporting.HistoryEntry> parent;
   com.googlecode.objectify.Key<? extends google.registry.model.billing.BillingEvent> cancellationMatchingBillingEvent;
+  com.googlecode.objectify.Key<google.registry.model.domain.token.AllocationToken> allocationToken;
   google.registry.model.billing.BillingEvent$Reason reason;
   java.lang.Integer periodYears;
   java.lang.String clientId;