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 f293ca262..4c8cda52a 100644 --- a/core/src/test/java/google/registry/flows/contact/ContactTransferRequestFlowTest.java +++ b/core/src/test/java/google/registry/flows/contact/ContactTransferRequestFlowTest.java @@ -20,13 +20,12 @@ import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.collect.MoreCollectors.onlyElement; import static com.google.common.truth.Truth.assertThat; import static google.registry.config.RegistryConfig.getContactAutomaticTransferLength; -import static google.registry.persistence.transaction.TransactionManagerFactory.tm; -import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm; import static google.registry.testing.ContactResourceSubject.assertAboutContacts; import static google.registry.testing.DatabaseHelper.assertNoBillingEvents; import static google.registry.testing.DatabaseHelper.assertPollMessagesEqual; import static google.registry.testing.DatabaseHelper.deleteResource; import static google.registry.testing.DatabaseHelper.getPollMessages; +import static google.registry.testing.DatabaseHelper.loadByKeys; import static google.registry.testing.DatabaseHelper.persistActiveContact; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions; @@ -136,10 +135,7 @@ class ContactTransferRequestFlowTest // poll messages, the approval notice ones for gaining and losing registrars. assertPollMessagesEqual( Iterables.filter( - transactIfJpaTm( - () -> tm().loadByKeys(contact.getTransferData().getServerApproveEntities())) - .values(), - PollMessage.class), + loadByKeys(contact.getTransferData().getServerApproveEntities()), PollMessage.class), ImmutableList.of(gainingApproveMessage, losingApproveMessage)); } 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 93c5b0daa..a5f4eb42d 100644 --- a/core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java +++ b/core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java @@ -35,6 +35,8 @@ import static google.registry.testing.DatabaseHelper.createTld; import static google.registry.testing.DatabaseHelper.getOnlyHistoryEntryOfType; import static google.registry.testing.DatabaseHelper.getOnlyPollMessage; import static google.registry.testing.DatabaseHelper.getPollMessages; +import static google.registry.testing.DatabaseHelper.loadByKey; +import static google.registry.testing.DatabaseHelper.loadByKeys; import static google.registry.testing.DatabaseHelper.loadRegistrar; import static google.registry.testing.DatabaseHelper.persistActiveContact; import static google.registry.testing.DatabaseHelper.persistResource; @@ -300,15 +302,13 @@ class DomainTransferRequestFlowTest // Assert that the domain's TransferData server-approve billing events match the above. if (expectTransferBillingEvent) { assertBillingEventsEqual( - transactIfJpaTm( - () -> tm().loadByKey(domain.getTransferData().getServerApproveBillingEvent())), + loadByKey(domain.getTransferData().getServerApproveBillingEvent()), optionalTransferBillingEvent.get()); } else { assertThat(domain.getTransferData().getServerApproveBillingEvent()).isNull(); } assertBillingEventsEqual( - transactIfJpaTm( - () -> tm().loadByKey(domain.getTransferData().getServerApproveAutorenewEvent())), + loadByKey(domain.getTransferData().getServerApproveAutorenewEvent()), gainingClientAutorenew); // Assert that the full set of server-approve billing events is exactly the extra ones plus // the transfer billing event (if present) and the gaining client autorenew. @@ -318,14 +318,10 @@ class DomainTransferRequestFlowTest .collect(toImmutableSet()); assertBillingEventsEqual( Iterables.filter( - transactIfJpaTm( - () -> - tm().loadByKeys(domain.getTransferData().getServerApproveEntities()).values()), - BillingEvent.class), + loadByKeys(domain.getTransferData().getServerApproveEntities()), BillingEvent.class), Sets.union(expectedServeApproveBillingEvents, extraBillingEvents)); // The domain's autorenew billing event should still point to the losing client's event. - BillingEvent.Recurring domainAutorenewEvent = - transactIfJpaTm(() -> tm().loadByKey(domain.getAutorenewBillingEvent())); + BillingEvent.Recurring domainAutorenewEvent = loadByKey(domain.getAutorenewBillingEvent()); assertThat(domainAutorenewEvent.getClientId()).isEqualTo("TheRegistrar"); assertThat(domainAutorenewEvent.getRecurrenceEndTime()).isEqualTo(implicitTransferTime); // The original grace periods should remain untouched. @@ -421,17 +417,13 @@ class DomainTransferRequestFlowTest // Assert that the poll messages show up in the TransferData server approve entities. assertPollMessagesEqual( - transactIfJpaTm( - () -> tm().loadByKey(domain.getTransferData().getServerApproveAutorenewPollMessage())), + loadByKey(domain.getTransferData().getServerApproveAutorenewPollMessage()), autorenewPollMessage); // Assert that the full set of server-approve poll messages is exactly the server approve // OneTime messages to gaining and losing registrars plus the gaining client autorenew. assertPollMessagesEqual( Iterables.filter( - transactIfJpaTm( - () -> - tm().loadByKeys(domain.getTransferData().getServerApproveEntities()).values()), - PollMessage.class), + loadByKeys(domain.getTransferData().getServerApproveEntities()), PollMessage.class), ImmutableList.of( transferApprovedPollMessage, losingTransferApprovedPollMessage, autorenewPollMessage)); } @@ -449,11 +441,7 @@ class DomainTransferRequestFlowTest .hasLastEppUpdateTime(implicitTransferTime) .and() .hasLastEppUpdateClientId("NewRegistrar"); - assertThat( - transactIfJpaTm( - () -> - tm().loadByKey(domainAfterAutomaticTransfer.getAutorenewBillingEvent()) - .getEventTime())) + assertThat(loadByKey(domainAfterAutomaticTransfer.getAutorenewBillingEvent()).getEventTime()) .isEqualTo(expectedExpirationTime); // And after the expected grace time, the grace period should be gone. DomainBase afterGracePeriod = diff --git a/core/src/test/java/google/registry/flows/host/HostUpdateFlowTest.java b/core/src/test/java/google/registry/flows/host/HostUpdateFlowTest.java index 6bda395ee..d072054ac 100644 --- a/core/src/test/java/google/registry/flows/host/HostUpdateFlowTest.java +++ b/core/src/test/java/google/registry/flows/host/HostUpdateFlowTest.java @@ -19,10 +19,10 @@ import static com.google.common.truth.Truth.assertThat; import static google.registry.batch.AsyncTaskEnqueuer.QUEUE_ASYNC_HOST_RENAME; import static google.registry.model.EppResourceUtils.loadByForeignKey; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; -import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm; import static google.registry.testing.DatabaseHelper.assertNoBillingEvents; import static google.registry.testing.DatabaseHelper.createTld; import static google.registry.testing.DatabaseHelper.getOnlyHistoryEntryOfType; +import static google.registry.testing.DatabaseHelper.loadByEntity; import static google.registry.testing.DatabaseHelper.newDomainBase; import static google.registry.testing.DatabaseHelper.newHostResource; import static google.registry.testing.DatabaseHelper.persistActiveDomain; @@ -302,8 +302,7 @@ class HostUpdateFlowTest extends ResourceFlowTestCase tm().loadByEntity(domain)).cloneProjectedAtTime(now); + DomainBase reloadedDomain = loadByEntity(domain).cloneProjectedAtTime(now); assertThat(reloadedDomain.getSubordinateHosts()).containsExactly("ns2.example.tld"); assertDnsTasksEnqueued("ns1.example.tld", "ns2.example.tld"); } @@ -337,15 +336,8 @@ class HostUpdateFlowTest extends ResourceFlowTestCase tm().loadByEntity(foo)) - .cloneProjectedAtTime(now) - .getSubordinateHosts()) - .isEmpty(); - assertThat( - transactIfJpaTm(() -> tm().loadByEntity(example)) - .cloneProjectedAtTime(now) - .getSubordinateHosts()) + assertThat(loadByEntity(foo).cloneProjectedAtTime(now).getSubordinateHosts()).isEmpty(); + assertThat(loadByEntity(example).cloneProjectedAtTime(now).getSubordinateHosts()) .containsExactly("ns2.example.tld"); assertDnsTasksEnqueued("ns2.foo.tld", "ns2.example.tld"); } @@ -380,11 +372,9 @@ class HostUpdateFlowTest extends ResourceFlowTestCase tm().loadByEntity(fooDomain)).cloneProjectedAtTime(now); + DomainBase reloadedFooDomain = loadByEntity(fooDomain).cloneProjectedAtTime(now); assertThat(reloadedFooDomain.getSubordinateHosts()).isEmpty(); - DomainBase reloadedTldDomain = - transactIfJpaTm(() -> tm().loadByEntity(tldDomain)).cloneProjectedAtTime(now); + DomainBase reloadedTldDomain = loadByEntity(tldDomain).cloneProjectedAtTime(now); assertThat(reloadedTldDomain.getSubordinateHosts()).containsExactly("ns2.example.tld"); assertDnsTasksEnqueued("ns1.example.foo", "ns2.example.tld"); } @@ -427,8 +417,7 @@ class HostUpdateFlowTest extends ResourceFlowTestCase tm().loadByEntity(domain)).cloneProjectedAtTime(clock.nowUtc()); + DomainBase reloadedDomain = loadByEntity(domain).cloneProjectedAtTime(clock.nowUtc()); assertThat(reloadedDomain.getSubordinateHosts()).isEmpty(); assertDnsTasksEnqueued("ns1.example.foo"); } @@ -464,10 +453,7 @@ class HostUpdateFlowTest extends ResourceFlowTestCase tm().loadByEntity(domain)) - .cloneProjectedAtTime(now) - .getSubordinateHosts()) + assertThat(loadByEntity(domain).cloneProjectedAtTime(now).getSubordinateHosts()) .containsExactly("ns2.example.tld"); assertDnsTasksEnqueued("ns2.example.tld"); } diff --git a/core/src/test/java/google/registry/model/billing/BillingEventTest.java b/core/src/test/java/google/registry/model/billing/BillingEventTest.java index 36ff24718..8b3b2f374 100644 --- a/core/src/test/java/google/registry/model/billing/BillingEventTest.java +++ b/core/src/test/java/google/registry/model/billing/BillingEventTest.java @@ -19,8 +19,9 @@ import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIM import static google.registry.model.ofy.ObjectifyService.ofy; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.persistence.transaction.TransactionManagerUtil.ofyTmOrDoNothing; -import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm; import static google.registry.testing.DatabaseHelper.createTld; +import static google.registry.testing.DatabaseHelper.loadByEntity; +import static google.registry.testing.DatabaseHelper.loadByKey; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.util.DateTimeUtils.END_OF_TIME; @@ -184,14 +185,11 @@ public class BillingEventTest extends EntityTestCase { @TestOfyAndSql void testPersistence() { - assertThat(transactIfJpaTm(() -> tm().loadByEntity(oneTime))).isEqualTo(oneTime); - assertThat(transactIfJpaTm(() -> tm().loadByEntity(oneTimeSynthetic))) - .isEqualTo(oneTimeSynthetic); - assertThat(transactIfJpaTm(() -> tm().loadByEntity(recurring))).isEqualTo(recurring); - assertThat(transactIfJpaTm(() -> tm().loadByEntity(cancellationOneTime))) - .isEqualTo(cancellationOneTime); - assertThat(transactIfJpaTm(() -> tm().loadByEntity(cancellationRecurring))) - .isEqualTo(cancellationRecurring); + assertThat(loadByEntity(oneTime)).isEqualTo(oneTime); + assertThat(loadByEntity(oneTimeSynthetic)).isEqualTo(oneTimeSynthetic); + assertThat(loadByEntity(recurring)).isEqualTo(recurring); + assertThat(loadByEntity(cancellationOneTime)).isEqualTo(cancellationOneTime); + assertThat(loadByEntity(cancellationRecurring)).isEqualTo(cancellationRecurring); ofyTmOrDoNothing(() -> assertThat(tm().loadByEntity(modification)).isEqualTo(modification)); } @@ -220,10 +218,8 @@ public class BillingEventTest extends EntityTestCase { @TestOfyAndSql void testCancellationMatching() { - VKey recurringKey = - transactIfJpaTm( - () -> tm().loadByEntity(oneTimeSynthetic).getCancellationMatchingBillingEvent()); - assertThat(transactIfJpaTm(() -> tm().loadByKey(recurringKey))).isEqualTo(recurring); + VKey recurringKey = loadByEntity(oneTimeSynthetic).getCancellationMatchingBillingEvent(); + assertThat(loadByKey(recurringKey)).isEqualTo(recurring); } @TestOfyOnly 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 fbbbe8241..5d1af30e4 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 @@ -22,9 +22,8 @@ import static google.registry.model.domain.token.AllocationToken.TokenStatus.NOT import static google.registry.model.domain.token.AllocationToken.TokenStatus.VALID; 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; -import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm; import static google.registry.testing.DatabaseHelper.createTld; +import static google.registry.testing.DatabaseHelper.loadByEntity; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.util.DateTimeUtils.START_OF_TIME; @@ -78,8 +77,7 @@ public class AllocationTokenTest extends EntityTestCase { .put(DateTime.now(UTC).plusWeeks(8), TokenStatus.ENDED) .build()) .build()); - assertThat(transactIfJpaTm(() -> tm().loadByEntity(unlimitedUseToken))) - .isEqualTo(unlimitedUseToken); + assertThat(loadByEntity(unlimitedUseToken)).isEqualTo(unlimitedUseToken); DomainBase domain = persistActiveDomain("example.foo"); Key historyEntryKey = Key.create(Key.create(domain), HistoryEntry.class, 1); @@ -92,7 +90,7 @@ public class AllocationTokenTest extends EntityTestCase { .setCreationTimeForTest(DateTime.parse("2010-11-12T05:00:00Z")) .setTokenType(SINGLE_USE) .build()); - assertThat(transactIfJpaTm(() -> tm().loadByEntity(singleUseToken))).isEqualTo(singleUseToken); + assertThat(loadByEntity(singleUseToken)).isEqualTo(singleUseToken); } @TestOfyOnly diff --git a/core/src/test/java/google/registry/testing/DatabaseHelper.java b/core/src/test/java/google/registry/testing/DatabaseHelper.java index 2102484ef..4bae9bab2 100644 --- a/core/src/test/java/google/registry/testing/DatabaseHelper.java +++ b/core/src/test/java/google/registry/testing/DatabaseHelper.java @@ -54,6 +54,7 @@ import static org.junit.jupiter.api.Assertions.fail; import com.google.common.base.Ascii; import com.google.common.base.Splitter; import com.google.common.base.Supplier; +import com.google.common.collect.ImmutableCollection; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -1244,5 +1245,45 @@ public class DatabaseHelper { clientId); } + /** + * Loads (i.e. reloads) the specified entity from the DB. + * + *

If the transaction manager is Cloud SQL, then this creates an inner wrapping transaction for + * convenience, so you don't need to wrap it in a transaction at the callsite. + */ + public static T loadByEntity(T entity) { + return transactIfJpaTm(() -> tm().loadByEntity(entity)); + } + + /** + * Loads the specified entity by its key from the DB. + * + *

If the transaction manager is Cloud SQL, then this creates an inner wrapping transaction for + * convenience, so you don't need to wrap it in a transaction at the callsite. + */ + public static T loadByKey(VKey key) { + return transactIfJpaTm(() -> tm().loadByKey(key)); + } + + /** + * Loads the specified entities by their keys from the DB. + * + *

If the transaction manager is Cloud SQL, then this creates an inner wrapping transaction for + * convenience, so you don't need to wrap it in a transaction at the callsite. + */ + public static ImmutableCollection loadByKeys(Iterable> keys) { + return transactIfJpaTm(() -> tm().loadByKeys(keys).values()); + } + + /** + * Loads all of the entities of the specified type from the DB. + * + *

If the transaction manager is Cloud SQL, then this creates an inner wrapping transaction for + * convenience, so you don't need to wrap it in a transaction at the callsite. + */ + public static ImmutableList loadAllOf(Class clazz) { + return transactIfJpaTm(() -> tm().loadAllOf(clazz)); + } + private DatabaseHelper() {} } diff --git a/core/src/test/java/google/registry/tools/DeleteAllocationTokensCommandTest.java b/core/src/test/java/google/registry/tools/DeleteAllocationTokensCommandTest.java index 9f3116681..bd1b56cb4 100644 --- a/core/src/test/java/google/registry/tools/DeleteAllocationTokensCommandTest.java +++ b/core/src/test/java/google/registry/tools/DeleteAllocationTokensCommandTest.java @@ -19,6 +19,7 @@ import static google.registry.model.domain.token.AllocationToken.TokenType.SINGL import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm; import static google.registry.testing.DatabaseHelper.createTlds; +import static google.registry.testing.DatabaseHelper.loadAllOf; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistResource; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -134,10 +135,9 @@ class DeleteAllocationTokensCommandTest extends CommandTestCase tm().loadAllOf(AllocationToken.class).size())).isEqualTo(56); + assertThat(loadAllOf(AllocationToken.class).size()).isEqualTo(56); runCommandForced("--prefix", "batch"); - assertThat(transactIfJpaTm(() -> tm().loadAllOf(AllocationToken.class).size())) - .isEqualTo(56 - 25); + assertThat(loadAllOf(AllocationToken.class).size()).isEqualTo(56 - 25); } @TestOfyAndSql diff --git a/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java b/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java index 73f564d30..98b52c0a5 100644 --- a/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java +++ b/core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java @@ -17,10 +17,9 @@ package google.registry.tools; import static com.google.common.truth.Truth.assertThat; 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; -import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm; import static google.registry.testing.DatabaseHelper.assertAllocationTokens; import static google.registry.testing.DatabaseHelper.createTld; +import static google.registry.testing.DatabaseHelper.loadAllOf; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.util.DateTimeUtils.START_OF_TIME; import static java.nio.charset.StandardCharsets.UTF_8; @@ -134,7 +133,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase tm().loadAllOf(AllocationToken.class).size())).isEqualTo(100); + assertThat(loadAllOf(AllocationToken.class).size()).isEqualTo(100); } @TestOfyAndSql @@ -199,7 +198,7 @@ class GenerateAllocationTokensCommandTest extends CommandTestCase sampleTokens = command.stringGenerator.createStrings(13, 100); runCommand("--tokens", Joiner.on(",").join(sampleTokens)); assertInStdout(Iterables.toArray(sampleTokens, String.class)); - assertThat(transactIfJpaTm(() -> tm().loadAllOf(AllocationToken.class).size())).isEqualTo(100); + assertThat(loadAllOf(AllocationToken.class).size()).isEqualTo(100); } @TestOfyAndSql