From d3b07f6ab0101f2ee80f7747a801c1baf2fee6ab Mon Sep 17 00:00:00 2001
From: gbrodman <gbrodman@google.com>
Date: Mon, 23 Aug 2021 13:08:09 -0600
Subject: [PATCH] Add double-replay to remaining existing ReplayExtension calls
 (#1297)

The only other change is that we need to reconstitute
serverApproveEntities for DomainTransferData in more situations (to fill
out the ofy keys)
---
 .../google/registry/model/transfer/DomainTransferData.java  | 6 ++++--
 .../batch/ExpandRecurringBillingEventsActionTest.java       | 2 +-
 .../google/registry/flows/contact/ContactCheckFlowTest.java | 2 +-
 .../registry/flows/contact/ContactCreateFlowTest.java       | 2 +-
 .../registry/flows/contact/ContactDeleteFlowTest.java       | 2 +-
 .../google/registry/flows/contact/ContactInfoFlowTest.java  | 2 +-
 .../flows/contact/ContactTransferApproveFlowTest.java       | 2 +-
 .../flows/contact/ContactTransferCancelFlowTest.java        | 2 +-
 .../flows/contact/ContactTransferQueryFlowTest.java         | 2 +-
 .../flows/contact/ContactTransferRejectFlowTest.java        | 2 +-
 .../flows/contact/ContactTransferRequestFlowTest.java       | 2 +-
 .../registry/flows/contact/ContactUpdateFlowTest.java       | 2 +-
 .../google/registry/flows/domain/DomainCheckFlowTest.java   | 2 +-
 .../registry/flows/domain/DomainClaimsCheckFlowTest.java    | 2 +-
 .../google/registry/flows/domain/DomainCreateFlowTest.java  | 2 +-
 .../google/registry/flows/domain/DomainDeleteFlowTest.java  | 2 +-
 .../google/registry/flows/domain/DomainInfoFlowTest.java    | 2 +-
 .../google/registry/flows/domain/DomainRenewFlowTest.java   | 2 +-
 .../registry/flows/domain/DomainRestoreRequestFlowTest.java | 2 +-
 .../flows/domain/DomainTransferApproveFlowTest.java         | 2 +-
 .../registry/flows/domain/DomainTransferCancelFlowTest.java | 2 +-
 .../registry/flows/domain/DomainTransferQueryFlowTest.java  | 2 +-
 .../registry/flows/domain/DomainTransferRejectFlowTest.java | 2 +-
 .../flows/domain/DomainTransferRequestFlowTest.java         | 2 +-
 .../google/registry/flows/domain/DomainUpdateFlowTest.java  | 2 +-
 .../java/google/registry/flows/poll/PollAckFlowTest.java    | 2 +-
 .../google/registry/flows/poll/PollRequestFlowTest.java     | 2 +-
 27 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/core/src/main/java/google/registry/model/transfer/DomainTransferData.java b/core/src/main/java/google/registry/model/transfer/DomainTransferData.java
index 6696bddd8..8c003e00f 100644
--- a/core/src/main/java/google/registry/model/transfer/DomainTransferData.java
+++ b/core/src/main/java/google/registry/model/transfer/DomainTransferData.java
@@ -157,10 +157,12 @@ public class DomainTransferData extends TransferData<DomainTransferData.Builder>
         DomainBase.restoreOfyFrom(rootKey, billingCancellationId, billingCancellationHistoryId);
 
     // Reconstruct server approve entities.  We currently have to call postLoad() a _second_ time
-    // if the billing cancellation id has been reconstituted, as it is part of that set.
+    // if any of the billing objects have been reconstituted, as they are part of that set.
     // TODO(b/183010623): Normalize the approaches to VKey reconstitution for the TransferData
     // hierarchy (the logic currently lives either in PostLoad or here, depending on the key).
-    if (billingCancellationId != null) {
+    if (billingCancellationId != null
+        || serverApproveBillingEvent != null
+        || serverApproveAutorenewEvent != null) {
       serverApproveEntities = null;
       postLoad();
     }
diff --git a/core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java b/core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java
index fcf4c4836..0ff7320dc 100644
--- a/core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java
+++ b/core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java
@@ -85,7 +85,7 @@ public class ExpandRecurringBillingEventsActionTest
 
   @Order(Order.DEFAULT - 2)
   @RegisterExtension
-  public final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  public final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   private DomainBase domain;
   private DomainHistory historyEntry;
diff --git a/core/src/test/java/google/registry/flows/contact/ContactCheckFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactCheckFlowTest.java
index cf7d1ad56..45fe8a8d9 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactCheckFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactCheckFlowTest.java
@@ -36,7 +36,7 @@ class ContactCheckFlowTest extends ResourceCheckFlowTestCase<ContactCheckFlow, C
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   ContactCheckFlowTest() {
     setEppInput("contact_check.xml");
diff --git a/core/src/test/java/google/registry/flows/contact/ContactCreateFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactCreateFlowTest.java
index 4823312b3..95a93a0b2 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactCreateFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactCreateFlowTest.java
@@ -45,7 +45,7 @@ class ContactCreateFlowTest extends ResourceFlowTestCase<ContactCreateFlow, Cont
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   ContactCreateFlowTest() {
     setEppInput("contact_create.xml");
diff --git a/core/src/test/java/google/registry/flows/contact/ContactDeleteFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactDeleteFlowTest.java
index 6d1a64fb2..f0d4360c9 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactDeleteFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactDeleteFlowTest.java
@@ -66,7 +66,7 @@ class ContactDeleteFlowTest extends ResourceFlowTestCase<ContactDeleteFlow, Cont
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void initFlowTest() {
diff --git a/core/src/test/java/google/registry/flows/contact/ContactInfoFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactInfoFlowTest.java
index ccce1cb89..9d7275f0d 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactInfoFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactInfoFlowTest.java
@@ -51,7 +51,7 @@ class ContactInfoFlowTest extends ResourceFlowTestCase<ContactInfoFlow, ContactR
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   ContactInfoFlowTest() {
     setEppInput("contact_info.xml");
diff --git a/core/src/test/java/google/registry/flows/contact/ContactTransferApproveFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactTransferApproveFlowTest.java
index 61615100f..f1054b813 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactTransferApproveFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactTransferApproveFlowTest.java
@@ -54,7 +54,7 @@ class ContactTransferApproveFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void setUp() {
diff --git a/core/src/test/java/google/registry/flows/contact/ContactTransferCancelFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactTransferCancelFlowTest.java
index 1f94f0572..6a456d097 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactTransferCancelFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactTransferCancelFlowTest.java
@@ -51,7 +51,7 @@ class ContactTransferCancelFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void setUp() {
diff --git a/core/src/test/java/google/registry/flows/contact/ContactTransferQueryFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactTransferQueryFlowTest.java
index 116a8bb09..ecf45f1ee 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactTransferQueryFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactTransferQueryFlowTest.java
@@ -46,7 +46,7 @@ class ContactTransferQueryFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void setUp() {
diff --git a/core/src/test/java/google/registry/flows/contact/ContactTransferRejectFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactTransferRejectFlowTest.java
index 22b3fd5ba..020edeb2f 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactTransferRejectFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactTransferRejectFlowTest.java
@@ -53,7 +53,7 @@ class ContactTransferRejectFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void setUp() {
diff --git a/core/src/test/java/google/registry/flows/contact/ContactTransferRequestFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactTransferRequestFlowTest.java
index fc81f0e70..dc7b304b6 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactTransferRequestFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactTransferRequestFlowTest.java
@@ -65,7 +65,7 @@ class ContactTransferRequestFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   ContactTransferRequestFlowTest() {
     // We need the transfer to happen at exactly this time in order for the response to match up.
diff --git a/core/src/test/java/google/registry/flows/contact/ContactUpdateFlowTest.java b/core/src/test/java/google/registry/flows/contact/ContactUpdateFlowTest.java
index db0c7087f..64df99434 100644
--- a/core/src/test/java/google/registry/flows/contact/ContactUpdateFlowTest.java
+++ b/core/src/test/java/google/registry/flows/contact/ContactUpdateFlowTest.java
@@ -53,7 +53,7 @@ class ContactUpdateFlowTest extends ResourceFlowTestCase<ContactUpdateFlow, Cont
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   ContactUpdateFlowTest() {
     setEppInput("contact_update.xml");
diff --git a/core/src/test/java/google/registry/flows/domain/DomainCheckFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainCheckFlowTest.java
index 1a00c6f86..420915266 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainCheckFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainCheckFlowTest.java
@@ -93,7 +93,7 @@ class DomainCheckFlowTest extends ResourceCheckFlowTestCase<DomainCheckFlow, Dom
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   DomainCheckFlowTest() {
     setEppInput("domain_check_one_tld.xml");
diff --git a/core/src/test/java/google/registry/flows/domain/DomainClaimsCheckFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainClaimsCheckFlowTest.java
index a480fc198..b66413512 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainClaimsCheckFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainClaimsCheckFlowTest.java
@@ -50,7 +50,7 @@ public class DomainClaimsCheckFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   DomainClaimsCheckFlowTest() {
     setEppInput("domain_check_claims.xml");
diff --git a/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java
index 10335c4ec..2fc2a8746 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java
@@ -188,7 +188,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   DomainCreateFlowTest() {
     setEppInput("domain_create.xml", ImmutableMap.of("DOMAIN", "example.tld"));
diff --git a/core/src/test/java/google/registry/flows/domain/DomainDeleteFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainDeleteFlowTest.java
index abb2e97f0..c6b650dd3 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainDeleteFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainDeleteFlowTest.java
@@ -119,7 +119,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   private DomainBase domain;
   private DomainHistory earlierHistoryEntry;
diff --git a/core/src/test/java/google/registry/flows/domain/DomainInfoFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainInfoFlowTest.java
index 63960e016..6fbae0726 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainInfoFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainInfoFlowTest.java
@@ -87,7 +87,7 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   /**
    * The domain_info_fee.xml default substitutions common to most tests.
diff --git a/core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java
index c183dfadb..dec969f8e 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java
@@ -112,7 +112,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void initDomainTest() {
diff --git a/core/src/test/java/google/registry/flows/domain/DomainRestoreRequestFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainRestoreRequestFlowTest.java
index aa4998883..46963848a 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainRestoreRequestFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainRestoreRequestFlowTest.java
@@ -90,7 +90,7 @@ class DomainRestoreRequestFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   private static final ImmutableMap<String, String> FEE_06_MAP =
       ImmutableMap.of("FEE_VERSION", "0.6", "FEE_NS", "fee", "CURRENCY", "USD");
diff --git a/core/src/test/java/google/registry/flows/domain/DomainTransferApproveFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainTransferApproveFlowTest.java
index e4fec138a..e5b76623c 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainTransferApproveFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainTransferApproveFlowTest.java
@@ -91,7 +91,7 @@ class DomainTransferApproveFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void setUp() {
diff --git a/core/src/test/java/google/registry/flows/domain/DomainTransferCancelFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainTransferCancelFlowTest.java
index 184208057..c40237b8a 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainTransferCancelFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainTransferCancelFlowTest.java
@@ -71,7 +71,7 @@ class DomainTransferCancelFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void setUp() {
diff --git a/core/src/test/java/google/registry/flows/domain/DomainTransferQueryFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainTransferQueryFlowTest.java
index 3901a4276..c7d8c61a4 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainTransferQueryFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainTransferQueryFlowTest.java
@@ -48,7 +48,7 @@ class DomainTransferQueryFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void beforeEach() {
diff --git a/core/src/test/java/google/registry/flows/domain/DomainTransferRejectFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainTransferRejectFlowTest.java
index 77d4c13ad..44e24798a 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainTransferRejectFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainTransferRejectFlowTest.java
@@ -73,7 +73,7 @@ class DomainTransferRejectFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void setUp() {
diff --git a/core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java
index eb8b7fae3..4d0fc8a78 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java
@@ -126,7 +126,7 @@ class DomainTransferRequestFlowTest
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   private static final ImmutableMap<String, String> BASE_FEE_MAP =
       new ImmutableMap.Builder<String, String>()
diff --git a/core/src/test/java/google/registry/flows/domain/DomainUpdateFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainUpdateFlowTest.java
index a2be1060e..12df3b101 100644
--- a/core/src/test/java/google/registry/flows/domain/DomainUpdateFlowTest.java
+++ b/core/src/test/java/google/registry/flows/domain/DomainUpdateFlowTest.java
@@ -120,7 +120,7 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   @BeforeEach
   void initDomainTest() {
diff --git a/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java b/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java
index 32843a4d8..0c9394994 100644
--- a/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java
+++ b/core/src/test/java/google/registry/flows/poll/PollAckFlowTest.java
@@ -51,7 +51,7 @@ class PollAckFlowTest extends FlowTestCase<PollAckFlow> {
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   /** This is the message id being sent in the ACK request. */
   private static final long MESSAGE_ID = 3;
diff --git a/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java b/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java
index 0856e37ab..df40b36b6 100644
--- a/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java
+++ b/core/src/test/java/google/registry/flows/poll/PollRequestFlowTest.java
@@ -58,7 +58,7 @@ class PollRequestFlowTest extends FlowTestCase<PollRequestFlow> {
 
   @Order(value = Order.DEFAULT - 2)
   @RegisterExtension
-  final ReplayExtension replayExtension = ReplayExtension.createWithCompare(clock);
+  final ReplayExtension replayExtension = ReplayExtension.createWithDoubleReplay(clock);
 
   private DomainBase domain;
   private ContactResource contact;