From f3d393661538e1363ddd08f3eef6cdae8b10addc Mon Sep 17 00:00:00 2001 From: gbrodman Date: Tue, 27 Jul 2021 16:05:59 -0400 Subject: [PATCH] Use the DB migration schedule for SQL->DS replay (#1242) This is instead of the current configuration parameter. In addition, this adds some helpers to DatabaseHelper to make the transitions easier, since we more frequently need to alter + reset the schedule. --- .../registry/config/RegistryConfig.java | 15 ------- .../config/RegistryConfigSettings.java | 1 - .../registry/config/files/default-config.yaml | 3 -- .../DatabaseMigrationStateSchedule.java | 2 +- .../JpaTransactionManagerImpl.java | 26 ++++++++++- .../ReplayCommitLogsToSqlActionTest.java | 3 +- .../DatabaseMigrationStateScheduleTest.java | 4 +- .../JpaTransactionManagerExtension.java | 2 + .../transaction/TransactionTest.java | 37 +++++++++++----- .../ReplicateToDatastoreActionTest.java | 29 +++--------- .../registry/testing/DatabaseHelper.java | 44 +++++++++++++++++++ .../registry/testing/ReplayExtension.java | 6 +-- .../GetDatabaseMigrationStateCommandTest.java | 3 +- .../SetDatabaseMigrationStateCommandTest.java | 16 +------ 14 files changed, 112 insertions(+), 79 deletions(-) diff --git a/core/src/main/java/google/registry/config/RegistryConfig.java b/core/src/main/java/google/registry/config/RegistryConfig.java index c98319435..5e19522a1 100644 --- a/core/src/main/java/google/registry/config/RegistryConfig.java +++ b/core/src/main/java/google/registry/config/RegistryConfig.java @@ -1494,21 +1494,6 @@ public final class RegistryConfig { return CONFIG_SETTINGS.get().hibernate.hikariIdleTimeout; } - /** - * Returns whether to replicate cloud SQL transactions to datastore. - * - *

If true, all cloud SQL transactions will be persisted as TransactionEntity objects in the - * Transaction table and replayed against datastore in a cron job. - */ - public static boolean getCloudSqlReplicateTransactions() { - return CONFIG_SETTINGS.get().cloudSql.replicateTransactions; - } - - @VisibleForTesting - public static void overrideCloudSqlReplicateTransactions(boolean replicateTransactions) { - CONFIG_SETTINGS.get().cloudSql.replicateTransactions = replicateTransactions; - } - /** Returns the roid suffix to be used for the roids of all contacts and hosts. */ public static String getContactAndHostRoidSuffix() { return CONFIG_SETTINGS.get().registryPolicy.contactAndHostRoidSuffix; diff --git a/core/src/main/java/google/registry/config/RegistryConfigSettings.java b/core/src/main/java/google/registry/config/RegistryConfigSettings.java index cf7c5ce22..4a07b658a 100644 --- a/core/src/main/java/google/registry/config/RegistryConfigSettings.java +++ b/core/src/main/java/google/registry/config/RegistryConfigSettings.java @@ -126,7 +126,6 @@ public class RegistryConfigSettings { // TODO(05012021): remove username field after it is removed from all yaml files. public String username; public String instanceConnectionName; - public boolean replicateTransactions; } /** Configuration for Apache Beam (Cloud Dataflow). */ diff --git a/core/src/main/java/google/registry/config/files/default-config.yaml b/core/src/main/java/google/registry/config/files/default-config.yaml index 729dd22b1..b61e3d600 100644 --- a/core/src/main/java/google/registry/config/files/default-config.yaml +++ b/core/src/main/java/google/registry/config/files/default-config.yaml @@ -227,9 +227,6 @@ cloudSql: jdbcUrl: jdbc:postgresql://localhost # This name is used by Cloud SQL when connecting to the database. instanceConnectionName: project-id:region:instance-id - # Set this to true to replicate cloud SQL transactions to datastore in the - # background. - replicateTransactions: false cloudDns: # Set both properties to null in Production. diff --git a/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java b/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java index 083f18557..46d30debc 100644 --- a/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java +++ b/core/src/main/java/google/registry/model/common/DatabaseMigrationStateSchedule.java @@ -225,7 +225,7 @@ public class DatabaseMigrationStateSchedule extends CrossTldSingleton @VisibleForTesting static TimedTransitionProperty getUncached() { return ofyTm() - .transact( + .transactNew( () -> ofyTm() .loadSingleton(DatabaseMigrationStateSchedule.class) diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java index 05f2e5ee3..c40db39cd 100644 --- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java +++ b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java @@ -30,8 +30,9 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Streams; import com.google.common.flogger.FluentLogger; -import google.registry.config.RegistryConfig; import google.registry.model.ImmutableObject; +import google.registry.model.common.DatabaseMigrationStateSchedule; +import google.registry.model.common.DatabaseMigrationStateSchedule.ReplayDirection; import google.registry.model.index.EppResourceIndex; import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex; import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex; @@ -103,6 +104,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager { // synchronously. private final ThreadLocal transactionInfo = ThreadLocal.withInitial(TransactionInfo::new); + // If this value is present, use it to determine whether or not to replay SQL transactions to + // Datastore, rather than using the schedule stored in Datastore. + private static ThreadLocal> replaySqlToDatastoreOverrideForTest = + ThreadLocal.withInitial(Optional::empty); public JpaTransactionManagerImpl(EntityManagerFactory emf, Clock clock) { this.emf = emf; @@ -736,6 +741,16 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager { return entity; } + /** Sets the override to always/never replay SQL transactions to Datastore. */ + public static void setReplaySqlToDatastoreOverrideForTest(boolean replaySqlToDs) { + replaySqlToDatastoreOverrideForTest.set(Optional.of(replaySqlToDs)); + } + + /** Removes the replay-SQL-to-Datastore override; the migration schedule will then be used. */ + public static void removeReplaySqlToDsOverrideForTest() { + replaySqlToDatastoreOverrideForTest.set(Optional.empty()); + } + private static class TransactionInfo { EntityManager entityManager; boolean inTransaction = false; @@ -755,7 +770,14 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager { checkArgumentNotNull(clock); inTransaction = true; transactionTime = clock.nowUtc(); - if (withBackup && RegistryConfig.getCloudSqlReplicateTransactions()) { + if (withBackup + && replaySqlToDatastoreOverrideForTest + .get() + .orElseGet( + () -> + DatabaseMigrationStateSchedule.getValueAtTime(transactionTime) + .getReplayDirection() + .equals(ReplayDirection.SQL_TO_DATASTORE))) { contentsBuilder = new Transaction.Builder(); } } diff --git a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java b/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java index 24c980084..7fe2998e1 100644 --- a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java +++ b/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java @@ -67,6 +67,7 @@ import google.registry.persistence.transaction.TransactionManagerFactory; import google.registry.schema.replay.SqlReplayCheckpoint; import google.registry.schema.tld.PremiumEntry; import google.registry.testing.AppEngineExtension; +import google.registry.testing.DatabaseHelper; import google.registry.testing.FakeClock; import google.registry.testing.FakeResponse; import google.registry.testing.TestObject; @@ -149,7 +150,7 @@ public class ReplayCommitLogsToSqlActionTest { @AfterEach void afterEach() { - ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + DatabaseHelper.removeDatabaseMigrationSchedule(); } @Test diff --git a/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java b/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java index 9906f21ea..e954ea608 100644 --- a/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java +++ b/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java @@ -15,7 +15,6 @@ package google.registry.model.common; import static com.google.common.truth.Truth.assertThat; -import static google.registry.model.common.DatabaseMigrationStateSchedule.DEFAULT_TRANSITION_MAP; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_ONLY; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.DATASTORE_PRIMARY_READ_ONLY; @@ -30,6 +29,7 @@ import static org.junit.Assert.assertThrows; import com.google.common.collect.ImmutableSortedMap; import google.registry.model.EntityTestCase; import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState; +import google.registry.testing.DatabaseHelper; import org.joda.time.DateTime; import org.joda.time.Duration; import org.junit.jupiter.api.AfterEach; @@ -45,7 +45,7 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase { @AfterEach void afterEach() { - ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + DatabaseHelper.removeDatabaseMigrationSchedule(); } @Test diff --git a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java index c4b649cf5..106fa1ee8 100644 --- a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java +++ b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java @@ -202,11 +202,13 @@ abstract class JpaTransactionManagerExtension implements BeforeEachCallback, Aft JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock); cachedTm = TransactionManagerFactory.jpaTm(); TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(txnManager)); + JpaTransactionManagerImpl.setReplaySqlToDatastoreOverrideForTest(false); } @Override public void afterEach(ExtensionContext context) { TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(cachedTm)); + JpaTransactionManagerImpl.removeReplaySqlToDsOverrideForTest(); cachedTm = null; } diff --git a/core/src/test/java/google/registry/persistence/transaction/TransactionTest.java b/core/src/test/java/google/registry/persistence/transaction/TransactionTest.java index 889f8b73a..c1215c7a2 100644 --- a/core/src/test/java/google/registry/persistence/transaction/TransactionTest.java +++ b/core/src/test/java/google/registry/persistence/transaction/TransactionTest.java @@ -22,36 +22,53 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import com.googlecode.objectify.Key; import com.googlecode.objectify.annotation.Entity; import com.googlecode.objectify.annotation.Id; -import google.registry.config.RegistryConfig; import google.registry.model.ImmutableObject; +import google.registry.model.ofy.Ofy; import google.registry.persistence.VKey; import google.registry.testing.AppEngineExtension; +import google.registry.testing.DatabaseHelper; +import google.registry.testing.FakeClock; +import google.registry.testing.InjectExtension; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.ObjectOutputStream; import java.io.StreamCorruptedException; +import org.joda.time.DateTime; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; class TransactionTest { + private final FakeClock fakeClock = new FakeClock(DateTime.parse("2000-01-01TZ")); + @RegisterExtension final AppEngineExtension appEngine = AppEngineExtension.builder() .withDatastoreAndCloudSql() + .withClock(fakeClock) .withOfyTestEntities(TestEntity.class) .withJpaUnitTestEntities(TestEntity.class) .build(); + @RegisterExtension public final InjectExtension inject = new InjectExtension(); + private TestEntity fooEntity, barEntity; @BeforeEach void beforeEach() { + JpaTransactionManagerImpl.removeReplaySqlToDsOverrideForTest(); + inject.setStaticField(Ofy.class, "clock", fakeClock); fooEntity = new TestEntity("foo"); barEntity = new TestEntity("bar"); } + @AfterEach + void afterEach() { + DatabaseHelper.removeDatabaseMigrationSchedule(); + } + @Test void testTransactionReplay() { Transaction txn = new Transaction.Builder().addUpdate(fooEntity).addUpdate(barEntity).build(); @@ -102,14 +119,13 @@ class TransactionTest { @Test void testTransactionSerialization() throws IOException { - RegistryConfig.overrideCloudSqlReplicateTransactions(true); - try { - jpaTm() - .transact( - () -> { - jpaTm().insert(fooEntity); - jpaTm().insert(barEntity); - }); + DatabaseHelper.setMigrationScheduleToSqlPrimary(fakeClock); + jpaTm() + .transact( + () -> { + jpaTm().insert(fooEntity); + jpaTm().insert(barEntity); + }); TransactionEntity txnEnt = jpaTm().transact(() -> jpaTm().loadByKey(VKey.createSql(TransactionEntity.class, 1L))); Transaction txn = Transaction.deserialize(txnEnt.getContents()); @@ -125,9 +141,6 @@ class TransactionTest { assertThat( jpaTm().transact(() -> jpaTm().exists(VKey.createSql(TransactionEntity.class, 2L)))) .isFalse(); - } finally { - RegistryConfig.overrideCloudSqlReplicateTransactions(false); - } } @Test diff --git a/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java b/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java index 6adf86b09..589747ac4 100644 --- a/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java +++ b/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java @@ -26,15 +26,16 @@ import com.google.common.testing.TestLogHandler; import com.googlecode.objectify.Key; import com.googlecode.objectify.annotation.Entity; import com.googlecode.objectify.annotation.Id; -import google.registry.config.RegistryConfig; import google.registry.model.ImmutableObject; import google.registry.model.common.DatabaseMigrationStateSchedule; import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState; import google.registry.model.ofy.CommitLogBucket; import google.registry.model.ofy.Ofy; import google.registry.persistence.VKey; +import google.registry.persistence.transaction.JpaTransactionManagerImpl; import google.registry.persistence.transaction.TransactionEntity; import google.registry.testing.AppEngineExtension; +import google.registry.testing.DatabaseHelper; import google.registry.testing.FakeClock; import google.registry.testing.InjectExtension; import java.util.List; @@ -68,41 +69,21 @@ public class ReplicateToDatastoreActionTest { @BeforeEach public void setUp() { + JpaTransactionManagerImpl.removeReplaySqlToDsOverrideForTest(); injectExtension.setStaticField(Ofy.class, "clock", fakeClock); // Use a single bucket to expose timestamp inversion problems. injectExtension.setStaticField( CommitLogBucket.class, "bucketIdSupplier", Suppliers.ofInstance(1)); fakeClock.setAutoIncrementByOneMilli(); - RegistryConfig.overrideCloudSqlReplicateTransactions(true); + DatabaseHelper.setMigrationScheduleToSqlPrimary(fakeClock); Logger.getLogger(ReplicateToDatastoreAction.class.getCanonicalName()).addHandler(logHandler); - DateTime now = fakeClock.nowUtc(); - ofyTm() - .transact( - () -> - DatabaseMigrationStateSchedule.set( - ImmutableSortedMap.of( - START_OF_TIME, - MigrationState.DATASTORE_ONLY, - now, - MigrationState.DATASTORE_PRIMARY, - now.plusHours(1), - MigrationState.DATASTORE_PRIMARY_READ_ONLY, - now.plusHours(2), - MigrationState.SQL_PRIMARY))); fakeClock.advanceBy(Duration.standardDays(1)); } @AfterEach public void tearDown() { + DatabaseHelper.removeDatabaseMigrationSchedule(); fakeClock.disableAutoIncrement(); - RegistryConfig.overrideCloudSqlReplicateTransactions(false); - ofyTm() - .transact( - () -> - ofyTm() - .loadSingleton(DatabaseMigrationStateSchedule.class) - .ifPresent(ofyTm()::delete)); - DatabaseMigrationStateSchedule.CACHE.invalidateAll(); } @RetryingTest(4) diff --git a/core/src/test/java/google/registry/testing/DatabaseHelper.java b/core/src/test/java/google/registry/testing/DatabaseHelper.java index 4acf05737..8ad8151d8 100644 --- a/core/src/test/java/google/registry/testing/DatabaseHelper.java +++ b/core/src/test/java/google/registry/testing/DatabaseHelper.java @@ -35,6 +35,7 @@ import static google.registry.model.ResourceTransferUtils.createTransferResponse import static google.registry.model.ofy.ObjectifyService.auditedOfy; import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY; import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.persistence.transaction.TransactionManagerUtil.ofyTmOrDoNothing; import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm; @@ -71,6 +72,8 @@ import google.registry.model.EppResourceUtils; import google.registry.model.billing.BillingEvent; import google.registry.model.billing.BillingEvent.Flag; import google.registry.model.billing.BillingEvent.Reason; +import google.registry.model.common.DatabaseMigrationStateSchedule; +import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState; import google.registry.model.contact.ContactAuthInfo; import google.registry.model.contact.ContactHistory; import google.registry.model.contact.ContactResource; @@ -123,6 +126,7 @@ import org.joda.money.CurrencyUnit; import org.joda.money.Money; import org.joda.time.DateTime; import org.joda.time.DateTimeZone; +import org.joda.time.Duration; /** Static utils for setting up test resources. */ public class DatabaseHelper { @@ -1360,5 +1364,45 @@ public class DatabaseHelper { return entity; } + /** + * Sets a SQL_PRIMARY state on the {@link DatabaseMigrationStateSchedule}. + * + *

In order to allow for tests to manipulate the clock how they need, we start the transitions + * one millisecond after the clock's current time (in case the clock's current value is + * START_OF_TIME). We then advance the clock one second so that we're in the SQL_PRIMARY phase. + * + *

We must use the current time, otherwise the setting of the migration state will fail due to + * an invalid transition. + */ + public static void setMigrationScheduleToSqlPrimary(FakeClock fakeClock) { + DateTime now = fakeClock.nowUtc(); + ofyTm() + .transact( + () -> + DatabaseMigrationStateSchedule.set( + ImmutableSortedMap.of( + START_OF_TIME, + MigrationState.DATASTORE_ONLY, + now.plusMillis(1), + MigrationState.DATASTORE_PRIMARY, + now.plusMillis(2), + MigrationState.DATASTORE_PRIMARY_READ_ONLY, + now.plusMillis(3), + MigrationState.SQL_PRIMARY))); + fakeClock.advanceBy(Duration.standardSeconds(1)); + } + + /** Removes the database migration schedule, in essence transitioning to DATASTORE_ONLY. */ + public static void removeDatabaseMigrationSchedule() { + // use the raw calls because going SQL_PRIMARY -> DATASTORE_ONLY is not valid + ofyTm() + .transact( + () -> + ofyTm() + .loadSingleton(DatabaseMigrationStateSchedule.class) + .ifPresent(ofyTm()::delete)); + DatabaseMigrationStateSchedule.CACHE.invalidateAll(); + } + private DatabaseHelper() {} } diff --git a/core/src/test/java/google/registry/testing/ReplayExtension.java b/core/src/test/java/google/registry/testing/ReplayExtension.java index 220e5aa5c..46671debd 100644 --- a/core/src/test/java/google/registry/testing/ReplayExtension.java +++ b/core/src/test/java/google/registry/testing/ReplayExtension.java @@ -22,12 +22,12 @@ import com.google.common.base.Suppliers; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.googlecode.objectify.Key; -import google.registry.config.RegistryConfig; import google.registry.model.ImmutableObject; import google.registry.model.ofy.CommitLogBucket; import google.registry.model.ofy.ReplayQueue; import google.registry.model.ofy.TransactionInfo; import google.registry.persistence.VKey; +import google.registry.persistence.transaction.JpaTransactionManagerImpl; import google.registry.persistence.transaction.TransactionEntity; import google.registry.schema.replay.DatastoreEntity; import google.registry.schema.replay.ReplicateToDatastoreAction; @@ -92,7 +92,7 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback { // transaction manager gets injected. inOfyContext = DualDatabaseTestInvocationContextProvider.inOfyContext(context); if (sqlToDsReplicator != null && !inOfyContext) { - RegistryConfig.overrideCloudSqlReplicateTransactions(true); + JpaTransactionManagerImpl.setReplaySqlToDatastoreOverrideForTest(true); } context.getStore(ExtensionContext.Namespace.GLOBAL).put(ReplayExtension.class, this); @@ -105,7 +105,7 @@ public class ReplayExtension implements BeforeEachCallback, AfterEachCallback { replay(); injectExtension.afterEach(context); if (sqlToDsReplicator != null) { - RegistryConfig.overrideCloudSqlReplicateTransactions(false); + JpaTransactionManagerImpl.setReplaySqlToDatastoreOverrideForTest(false); } } diff --git a/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java index 5698c110c..ab171558b 100644 --- a/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java +++ b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java @@ -21,6 +21,7 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME; import com.google.common.collect.ImmutableSortedMap; import google.registry.model.common.DatabaseMigrationStateSchedule; import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState; +import google.registry.testing.DatabaseHelper; import google.registry.testing.DualDatabaseTest; import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; @@ -33,7 +34,7 @@ public class GetDatabaseMigrationStateCommandTest @AfterEach void afterEach() { - ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + DatabaseHelper.removeDatabaseMigrationSchedule(); } @TestOfyAndSql diff --git a/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java index 5781e8998..9b701dc5b 100644 --- a/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java +++ b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java @@ -25,32 +25,20 @@ import com.beust.jcommander.ParameterException; import com.google.common.collect.ImmutableSortedMap; import google.registry.model.common.DatabaseMigrationStateSchedule; import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState; +import google.registry.testing.DatabaseHelper; import google.registry.testing.DualDatabaseTest; import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; /** Tests for {@link SetDatabaseMigrationStateCommand}. */ @DualDatabaseTest public class SetDatabaseMigrationStateCommandTest extends CommandTestCase { - @BeforeEach - void beforeEach() { - // clear out any static state that may have been persisted - ofyTm() - .transact( - () -> - ofyTm() - .loadSingleton(DatabaseMigrationStateSchedule.class) - .ifPresent(ofyTm()::delete)); - DatabaseMigrationStateSchedule.CACHE.invalidateAll(); - } - @AfterEach void afterEach() { - ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + DatabaseHelper.removeDatabaseMigrationSchedule(); } @TestOfyAndSql