diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java index 8ed871803..bdeee15e1 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java +++ b/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java @@ -16,17 +16,22 @@ package google.registry.persistence.transaction; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; +import static org.joda.time.DateTimeZone.UTC; import com.google.appengine.api.utils.SystemProperty; import com.google.appengine.api.utils.SystemProperty.Environment.Value; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Suppliers; import google.registry.config.RegistryEnvironment; +import google.registry.model.common.DatabaseMigrationStateSchedule; +import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase; import google.registry.model.ofy.DatastoreTransactionManager; import google.registry.persistence.DaggerPersistenceComponent; import google.registry.tools.RegistryToolEnvironment; import google.registry.util.NonFinalForTesting; +import java.util.Optional; import java.util.function.Supplier; +import org.joda.time.DateTime; /** Factory class to create {@link TransactionManager} instance. */ // TODO: Rename this to PersistenceFactory and move to persistence package. @@ -34,7 +39,8 @@ public class TransactionManagerFactory { private static final DatastoreTransactionManager ofyTm = createTransactionManager(); - @NonFinalForTesting private static TransactionManager tm = ofyTm; + /** Optional override to manually set the transaction manager per-test. */ + private static Optional tmForTest = Optional.empty(); /** Supplier for jpaTm so that it is initialized only once, upon first usage. */ @NonFinalForTesting @@ -70,9 +76,19 @@ public class TransactionManagerFactory { return SystemProperty.environment.value() == Value.Production; } - /** Returns {@link TransactionManager} instance. */ + /** + * Returns the {@link TransactionManager} instance. + * + *

Returns the {@link JpaTransactionManager} or {@link DatastoreTransactionManager} based on + * the migration schedule or the manually specified per-test transaction manager. + */ public static TransactionManager tm() { - return tm; + if (tmForTest.isPresent()) { + return tmForTest.get(); + } + PrimaryDatabase primaryDatabase = + DatabaseMigrationStateSchedule.getValueAtTime(DateTime.now(UTC)).getPrimaryDatabase(); + return primaryDatabase.equals(PrimaryDatabase.DATASTORE) ? ofyTm() : jpaTm(); } /** @@ -111,9 +127,18 @@ public class TransactionManagerFactory { jpaTm = Suppliers.memoize(jpaTmSupplier::get); } - /** Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. */ + /** + * Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. + * + *

Used when overriding the per-test transaction manager for dual-database tests. + */ @VisibleForTesting - public static void setTm(TransactionManager newTm) { - tm = newTm; + public static void setTmForTest(TransactionManager newTm) { + tmForTest = Optional.of(newTm); + } + + /** Resets the overridden transaction manager post-test. */ + public static void removeTmOverrideForTest() { + tmForTest = Optional.empty(); } } diff --git a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java b/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java index 40c7a2db1..6ba146ba9 100644 --- a/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java +++ b/core/src/test/java/google/registry/backup/ReplayCommitLogsToSqlActionTest.java @@ -74,6 +74,7 @@ import google.registry.util.RequestStatusChecker; import java.io.IOException; import org.joda.time.DateTime; import org.joda.time.Duration; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -144,6 +145,11 @@ public class ReplayCommitLogsToSqlActionTest { TestObject.beforeSqlDeleteCallCount = 0; } + @AfterEach + void afterEach() { + ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + } + @Test void testReplay_multipleDiffFiles() throws Exception { jpaTm() diff --git a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java index 2610ef685..eba0ed1c6 100644 --- a/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java +++ b/core/src/test/java/google/registry/beam/rde/RdePipelineTest.java @@ -23,7 +23,7 @@ import static google.registry.model.common.Cursor.CursorType.RDE_STAGING; import static google.registry.model.rde.RdeMode.FULL; import static google.registry.model.rde.RdeMode.THIN; import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm; -import static google.registry.persistence.transaction.TransactionManagerFactory.setTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.rde.RdeResourceType.CONTACT; import static google.registry.rde.RdeResourceType.DOMAIN; @@ -122,7 +122,7 @@ public class RdePipelineTest { @BeforeEach void beforeEach() throws Exception { originalTm = tm(); - setTm(jpaTm()); + setTmForTest(jpaTm()); loadInitialData(); // Two real registrars have been created by loadInitialData(), named "New Registrar" and "The @@ -169,7 +169,7 @@ public class RdePipelineTest { @AfterEach void afterEach() { - setTm(originalTm); + setTmForTest(originalTm); } @Test diff --git a/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java b/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java index 08f015649..1a9f92a61 100644 --- a/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java +++ b/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java @@ -18,7 +18,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence; import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm; -import static google.registry.persistence.transaction.TransactionManagerFactory.setTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.setTmForTest; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.testing.AppEngineExtension.makeRegistrar1; import static google.registry.testing.DatabaseHelper.createTld; @@ -231,7 +231,7 @@ class Spec11PipelineTest { private void setupCloudSql() { TransactionManager originalTm = tm(); - setTm(jpaTm()); + setTmForTest(jpaTm()); persistNewRegistrar("TheRegistrar"); persistNewRegistrar("NewRegistrar"); Registrar registrar1 = @@ -271,7 +271,7 @@ class Spec11PipelineTest { persistResource(createDomain("no-email.com", "2A4BA9BBC-COM", registrar2, contact2)); persistResource( createDomain("anti-anti-anti-virus.dev", "555666888-DEV", registrar3, contact3)); - setTm(originalTm); + setTmForTest(originalTm); } private void verifySaveToGcs() throws Exception { 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 34b00a0e2..9906f21ea 100644 --- a/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java +++ b/core/src/test/java/google/registry/model/common/DatabaseMigrationStateScheduleTest.java @@ -15,6 +15,7 @@ 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; @@ -22,6 +23,7 @@ import static google.registry.model.common.DatabaseMigrationStateSchedule.Migrat import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY; import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState.SQL_PRIMARY_READ_ONLY; import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.util.DateTimeUtils.START_OF_TIME; import static org.junit.Assert.assertThrows; @@ -30,6 +32,7 @@ import google.registry.model.EntityTestCase; import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState; import org.joda.time.DateTime; import org.joda.time.Duration; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -40,6 +43,11 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase { fakeClock.setAutoIncrementByOneMilli(); } + @AfterEach + void afterEach() { + ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + } + @Test void testEmpty_returnsDatastoreOnlyMap() { assertThat(DatabaseMigrationStateSchedule.getUncached()) @@ -134,6 +142,15 @@ public class DatabaseMigrationStateScheduleTest extends EntityTestCase { .isEqualTo("Must be called in a transaction"); } + @Test + void testSuccess_factoryUsesSchedule() { + assertThat(tm().isOfy()).isTrue(); + // set the schedule to have converted to SQL_PRIMARY in the past + fakeClock.setTo(START_OF_TIME.plusDays(1)); + runValidTransition(DATASTORE_PRIMARY_READ_ONLY, SQL_PRIMARY); + assertThat(tm().isOfy()).isFalse(); + } + private void runValidTransition(MigrationState from, MigrationState to) { ImmutableSortedMap transitions = createMapEndingWithTransition(from, to); 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 e24e5cee8..6adf86b09 100644 --- a/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java +++ b/core/src/test/java/google/registry/schema/replay/ReplicateToDatastoreActionTest.java @@ -96,6 +96,13 @@ public class ReplicateToDatastoreActionTest { public void tearDown() { 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/DualDatabaseTestInvocationContextProvider.java b/core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java index 05a3b4296..2c3413f85 100644 --- a/core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java +++ b/core/src/test/java/google/registry/testing/DualDatabaseTestInvocationContextProvider.java @@ -154,15 +154,13 @@ class DualDatabaseTestInvocationContextProvider implements TestTemplateInvocatio context.getStore(NAMESPACE).put(ORIGINAL_TM_KEY, tm()); DatabaseType databaseType = (DatabaseType) context.getStore(NAMESPACE).get(INJECTED_TM_SUPPLIER_KEY); - TransactionManagerFactory.setTm(databaseType.getTm()); + TransactionManagerFactory.setTmForTest(databaseType.getTm()); } } static void restoreTmAfterDualDatabaseTest(ExtensionContext context) { if (isDualDatabaseTest(context)) { - TransactionManager original = - (TransactionManager) context.getStore(NAMESPACE).get(ORIGINAL_TM_KEY); - TransactionManagerFactory.setTm(original); + TransactionManagerFactory.removeTmOverrideForTest(); } } diff --git a/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java index 37a479c1e..5698c110c 100644 --- a/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java +++ b/core/src/test/java/google/registry/tools/GetDatabaseMigrationStateCommandTest.java @@ -24,15 +24,15 @@ import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStat import google.registry.testing.DualDatabaseTest; import google.registry.testing.TestOfyAndSql; import org.joda.time.DateTime; -import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.AfterEach; /** Tests for {@link GetDatabaseMigrationStateCommand}. */ @DualDatabaseTest public class GetDatabaseMigrationStateCommandTest extends CommandTestCase { - @BeforeEach - void beforeEach() { + @AfterEach + void afterEach() { ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); } diff --git a/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java index 2be0c8890..5781e8998 100644 --- a/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java +++ b/core/src/test/java/google/registry/tools/SetDatabaseMigrationStateCommandTest.java @@ -28,6 +28,7 @@ import google.registry.model.common.DatabaseMigrationStateSchedule.MigrationStat 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}. */ @@ -47,6 +48,11 @@ public class SetDatabaseMigrationStateCommandTest DatabaseMigrationStateSchedule.CACHE.invalidateAll(); } + @AfterEach + void afterEach() { + ofyTm().transact(() -> DatabaseMigrationStateSchedule.set(DEFAULT_TRANSITION_MAP.toValueMap())); + } + @TestOfyAndSql void testSuccess_setsBasicSchedule() throws Exception { assertThat(DatabaseMigrationStateSchedule.get()).isEqualTo(DEFAULT_TRANSITION_MAP);