Improve and consolidate testing of poll messages/billing events (#745)

* Improve and consolidate testing of poll messages/billing events

core/src/test/java/google/registry/flows/FlowTestCase.java

@@ -16,15 +16,12 @@ package google.registry.flows;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.model.eppcommon.EppXmlTransformer.marshal;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.testing.DatastoreHelper.POLL_MESSAGE_ID_STRIPPER;
-import static google.registry.testing.DatastoreHelper.getPollMessages;
 import static google.registry.testing.DatastoreHelper.stripBillingEventId;
 import static google.registry.xml.XmlTestUtils.assertXmlEquals;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -34,7 +31,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 import com.google.common.collect.ObjectArrays;
-import com.google.common.collect.Streams;
 import google.registry.config.RegistryConfig.ConfigModule.TmchCaMode;
 import google.registry.flows.EppTestComponent.FakesAndMocksModule;
 import google.registry.flows.picker.FlowPicker;
@@ -44,7 +40,6 @@ import google.registry.model.eppcommon.ProtocolDefinition;
 import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.ofy.Ofy;
-import google.registry.model.poll.PollMessage;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.tmch.ClaimsListShard.ClaimsListSingleton;
 import google.registry.monitoring.whitebox.EppMetric;
@@ -214,24 +209,6 @@ public abstract class FlowTestCase<F extends Flow> {
         .isEqualTo(canonicalizeGracePeriods(expected));
   }
 
-  protected void assertPollMessages(String clientId, PollMessage... expected) {
-    assertPollMessagesHelper(getPollMessages(clientId), expected);
-  }
-
-  protected void assertPollMessages(PollMessage... expected) {
-    assertPollMessagesHelper(getPollMessages(), expected);
-  }
-
-  /** Assert that the list matches all the poll messages in the fake Datastore. */
-  private void assertPollMessagesHelper(
-      Iterable<PollMessage> pollMessages, PollMessage... expected) {
-    // Ordering is irrelevant but duplicates should be considered independently.
-    assertThat(
-            Streams.stream(pollMessages).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Arrays.stream(expected).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()));
-  }
-
   private EppOutput runFlowInternal(CommitMode commitMode, UserPrivileges userPrivileges)
       throws Exception {
     eppMetricBuilder = EppMetric.builderForRequest(clock);

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

@@ -36,6 +36,7 @@ import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_DELETE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
 import static google.registry.testing.DatastoreHelper.getOnlyPollMessage;

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

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.flows.domain.DomainTransferFlowTestCase.persistWithPendingTransfer;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
 import static google.registry.testing.DatastoreHelper.loadRegistrar;

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

@@ -17,6 +17,7 @@ package google.registry.flows.domain;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
 import static google.registry.testing.DatastoreHelper.getPollMessages;

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

@@ -21,6 +21,7 @@ import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_CANCEL;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
+import static google.registry.testing.DatastoreHelper.assertPollMessages;
 import static google.registry.testing.DatastoreHelper.createPollMessageForImplicitTransfer;
 import static google.registry.testing.DatastoreHelper.deleteResource;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;

core/src/test/java/google/registry/testing/DatastoreHelper.java

@@ -26,6 +26,7 @@ import static google.registry.config.RegistryConfig.getContactAndHostRoidSuffix;
 import static google.registry.config.RegistryConfig.getContactAutomaticTransferLength;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.model.ResourceTransferUtils.createTransferResponse;
 import static google.registry.model.ofy.ObjectifyService.ofy;
@@ -52,7 +53,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Streams;
 import com.google.common.net.InetAddresses;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.cmd.Saver;
@@ -102,7 +102,6 @@ import google.registry.tmch.LordnTaskUtils;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
-import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
@@ -688,20 +687,15 @@ public class DatastoreHelper {
 
   /** Assert that the actual billing event matches the expected one, ignoring IDs. */
   public static void assertBillingEventsEqual(BillingEvent actual, BillingEvent expected) {
-    assertThat(stripBillingEventId(actual)).isEqualTo(stripBillingEventId(expected));
+    assertAboutImmutableObjects().that(actual).isEqualExceptFields(expected, "id");
   }
 
   /** Assert that the actual billing events match the expected ones, ignoring IDs and order. */
   public static void assertBillingEventsEqual(
       Iterable<BillingEvent> actual, Iterable<BillingEvent> expected) {
-    assertThat(
+    assertThat(actual)
-            Streams.stream(actual)
+        .comparingElementsUsing(immutableObjectCorrespondence("id"))
-                .map(DatastoreHelper::stripBillingEventId)
+        .containsExactlyElementsIn(expected);
-                .collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Streams.stream(expected)
-                .map(DatastoreHelper::stripBillingEventId)
-                .collect(toImmutableList()));
   }
 
   /** Assert that the expected billing events are exactly the ones found in the fake Datastore. */
@@ -719,14 +713,7 @@ public class DatastoreHelper {
    */
   public static void assertBillingEventsForResource(
       EppResource resource, BillingEvent... expected) {
-    assertThat(
+    assertBillingEventsEqual(getBillingEvents(resource), Arrays.asList(expected));
-            Streams.stream(getBillingEvents(resource))
-                .map(DatastoreHelper::stripBillingEventId)
-                .collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Arrays.stream(expected)
-                .map(DatastoreHelper::stripBillingEventId)
-                .collect(toImmutableList()));
   }
 
   /** Assert that there are no billing events. */
@@ -734,74 +721,71 @@ public class DatastoreHelper {
     assertThat(getBillingEvents()).isEmpty();
   }
 
-  /** Strips the billing event ID (really, sets it to a constant value) to facilitate comparison. */
+  /**
+   * Strips the billing event ID (really, sets it to a constant value) to facilitate comparison.
+   *
+   * <p>Note: Prefer {@link #assertPollMessagesEqual} when that is suitable.
+   */
   public static BillingEvent stripBillingEventId(BillingEvent billingEvent) {
     return billingEvent.asBuilder().setId(1L).build();
   }
 
   /** Assert that the actual poll message matches the expected one, ignoring IDs. */
   public static void assertPollMessagesEqual(PollMessage actual, PollMessage expected) {
-    assertThat(POLL_MESSAGE_ID_STRIPPER.apply(actual))
+    assertAboutImmutableObjects().that(actual).isEqualExceptFields(expected, "id");
-        .isEqualTo(POLL_MESSAGE_ID_STRIPPER.apply(expected));
   }
 
   /** Assert that the actual poll messages match the expected ones, ignoring IDs and order. */
   public static void assertPollMessagesEqual(
       Iterable<PollMessage> actual, Iterable<PollMessage> expected) {
-    assertThat(Streams.stream(actual).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()))
+    assertThat(actual)
-        .containsExactlyElementsIn(
+        .comparingElementsUsing(immutableObjectCorrespondence("id"))
-            Streams.stream(expected).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()));
+        .containsExactlyElementsIn(expected);
+  }
+
+  public static void assertPollMessages(String clientId, PollMessage... expected) {
+    assertPollMessagesEqual(getPollMessages(clientId), Arrays.asList(expected));
+  }
+
+  public static void assertPollMessages(PollMessage... expected) {
+    assertPollMessagesEqual(getPollMessages(), Arrays.asList(expected));
   }
 
   public static void assertPollMessagesForResource(EppResource resource, PollMessage... expected) {
-    assertThat(
+    assertPollMessagesEqual(getPollMessages(resource), Arrays.asList(expected));
-            getPollMessages(resource)
-                .stream()
-                .map(POLL_MESSAGE_ID_STRIPPER)
-                .collect(toImmutableList()))
-        .containsExactlyElementsIn(
-            Arrays.stream(expected).map(POLL_MESSAGE_ID_STRIPPER).collect(toImmutableList()));
   }
 
-  /** Helper to effectively erase the poll message ID to facilitate comparison. */
-  public static final Function<PollMessage, PollMessage> POLL_MESSAGE_ID_STRIPPER =
-      pollMessage -> pollMessage.asBuilder().setId(1L).build();
-
   public static ImmutableList<PollMessage> getPollMessages() {
-    return Streams.stream(ofy().load().type(PollMessage.class)).collect(toImmutableList());
+    return ImmutableList.copyOf(ofy().load().type(PollMessage.class));
   }
 
   public static ImmutableList<PollMessage> getPollMessages(String clientId) {
-    return Streams.stream(ofy().load().type(PollMessage.class).filter("clientId", clientId))
+    return ImmutableList.copyOf(ofy().load().type(PollMessage.class).filter("clientId", clientId));
-        .collect(toImmutableList());
   }
 
   public static ImmutableList<PollMessage> getPollMessages(EppResource resource) {
-    return Streams.stream(ofy().load().type(PollMessage.class).ancestor(resource))
+    return ImmutableList.copyOf(ofy().load().type(PollMessage.class).ancestor(resource));
-        .collect(toImmutableList());
   }
 
   public static ImmutableList<PollMessage> getPollMessages(String clientId, DateTime now) {
-    return Streams.stream(
+    return ImmutableList.copyOf(
         ofy()
             .load()
             .type(PollMessage.class)
             .filter("clientId", clientId)
-                .filter("eventTime <=", now.toDate()))
+            .filter("eventTime <=", now.toDate()));
-        .collect(toImmutableList());
   }
 
   /** Gets all PollMessages associated with the given EppResource. */
   public static ImmutableList<PollMessage> getPollMessages(
       EppResource resource, String clientId, DateTime now) {
-    return Streams.stream(
+    return ImmutableList.copyOf(
         ofy()
             .load()
             .type(PollMessage.class)
             .ancestor(resource)
             .filter("clientId", clientId)
-                .filter("eventTime <=", now.toDate()))
+            .filter("eventTime <=", now.toDate()));
-        .collect(toImmutableList());
   }
 
   public static PollMessage getOnlyPollMessage(String clientId) {