From cd6cb10b3767fed408c0a59e716462ccaac88f53 Mon Sep 17 00:00:00 2001 From: Ben McIlwain Date: Thu, 25 Mar 2021 16:14:46 -0400 Subject: [PATCH] Add some load convenience methods to DatabaseHelper (#1038) * Add some load convenience methods to DatabaseHelper These can only be called by test code, and they automatically wrap the load in a transaction if one isn't already specified (for convenience). In production code we don't want to be able to use these, as we have to be more thoughtful about transactions in production code (e.g. make sure that we aren't loading and then saving a resource in separate transactions in a way that makes it prone to contention errors). --- .../ContactTransferRequestFlowTest.java | 8 +--- .../domain/DomainTransferRequestFlowTest.java | 30 ++++---------- .../flows/host/HostUpdateFlowTest.java | 30 ++++---------- .../model/billing/BillingEventTest.java | 22 ++++------ .../domain/token/AllocationTokenTest.java | 8 ++-- .../registry/testing/DatabaseHelper.java | 41 +++++++++++++++++++ .../DeleteAllocationTokensCommandTest.java | 6 +-- .../GenerateAllocationTokensCommandTest.java | 7 ++-- 8 files changed, 78 insertions(+), 74 deletions(-) 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