Fix flaky test issues caused by lack of ofy init (#1246)

core/src/test/java/google/registry/beam/rde/RdePipelineTest.java

@@ -23,7 +23,6 @@ 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.setTmForTest;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.rde.RdeResourceType.CONTACT;
 import static google.registry.rde.RdeResourceType.DOMAIN;
@@ -50,7 +49,7 @@ import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
 import google.registry.persistence.transaction.JpaTestRules;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestExtension;
-import google.registry.persistence.transaction.TransactionManager;
+import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.rde.DepositFragment;
 import google.registry.rde.PendingDeposit;
 import google.registry.rde.RdeResourceType;
@@ -117,12 +116,9 @@ public class RdePipelineTest {
 
   private RdePipeline rdePipeline;
 
-  private TransactionManager originalTm;
-
   @BeforeEach
   void beforeEach() throws Exception {
-    originalTm = tm();
-    setTmForTest(jpaTm());
+    TransactionManagerFactory.setTmForTest(jpaTm());
     loadInitialData();
 
     // Two real registrars have been created by loadInitialData(), named "New Registrar" and "The
@@ -169,7 +165,7 @@ public class RdePipelineTest {
 
   @AfterEach
   void afterEach() {
-    setTmForTest(originalTm);
+    TransactionManagerFactory.removeTmOverrideForTest();
   }
 
   @Test

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java

@@ -17,7 +17,7 @@ package google.registry.persistence.transaction;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
-import static google.registry.testing.DatabaseHelper.assertDetached;
+import static google.registry.testing.DatabaseHelper.assertDetachedFromEntityManager;
 import static google.registry.testing.TestDataHelper.fileClassPath;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
@@ -388,7 +388,8 @@ class JpaTransactionManagerImplTest {
   void load_succeeds() {
     assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().insert(theEntity));
-    TestEntity persisted = jpaTm().transact(() -> assertDetached(jpaTm().loadByKey(theEntityKey)));
+    TestEntity persisted =
+        jpaTm().transact(() -> assertDetachedFromEntityManager(jpaTm().loadByKey(theEntityKey)));
     assertThat(persisted.name).isEqualTo("theEntity");
     assertThat(persisted.data).isEqualTo("foo");
   }
@@ -404,7 +405,8 @@ class JpaTransactionManagerImplTest {
   @Test
   void loadByEntity_succeeds() {
     jpaTm().transact(() -> jpaTm().insert(theEntity));
-    TestEntity persisted = jpaTm().transact(() -> assertDetached(jpaTm().loadByEntity(theEntity)));
+    TestEntity persisted =
+        jpaTm().transact(() -> assertDetachedFromEntityManager(jpaTm().loadByEntity(theEntity)));
     assertThat(persisted.name).isEqualTo("theEntity");
     assertThat(persisted.data).isEqualTo("foo");
   }
@@ -414,7 +416,11 @@ class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().exists(theEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().insert(theEntity));
     TestEntity persisted =
-        jpaTm().transact(() -> assertDetached(jpaTm().loadByKeyIfPresent(theEntityKey).get()));
+        jpaTm()
+            .transact(
+                () ->
+                    assertDetachedFromEntityManager(
+                        jpaTm().loadByKeyIfPresent(theEntityKey).get()));
     assertThat(persisted.name).isEqualTo("theEntity");
     assertThat(persisted.data).isEqualTo("foo");
   }
@@ -431,7 +437,9 @@ class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().exists(compoundIdEntity))).isFalse();
     jpaTm().transact(() -> jpaTm().insert(compoundIdEntity));
     TestCompoundIdEntity persisted =
-        jpaTm().transact(() -> assertDetached(jpaTm().loadByKey(compoundIdEntityKey)));
+        jpaTm()
+            .transact(
+                () -> assertDetachedFromEntityManager(jpaTm().loadByKey(compoundIdEntityKey)));
     assertThat(persisted.name).isEqualTo("compoundIdEntity");
     assertThat(persisted.age).isEqualTo(10);
     assertThat(persisted.data).isEqualTo("foo");
@@ -450,7 +458,7 @@ class JpaTransactionManagerImplTest {
                               theEntityKey, VKey.createSql(TestEntity.class, "does-not-exist")));
 
               assertThat(results).containsExactly(theEntityKey, theEntity);
-              assertDetached(results.get(theEntityKey));
+              assertDetachedFromEntityManager(results.get(theEntityKey));
             });
   }
 
@@ -463,7 +471,7 @@ class JpaTransactionManagerImplTest {
               ImmutableMap<VKey<? extends TestEntity>, TestEntity> results =
                   jpaTm().loadByKeysIfPresent(ImmutableList.of(theEntityKey));
               assertThat(results).containsExactly(theEntityKey, theEntity);
-              assertDetached(results.get(theEntityKey));
+              assertDetachedFromEntityManager(results.get(theEntityKey));
             });
   }
 
@@ -478,7 +486,7 @@ class JpaTransactionManagerImplTest {
                       .loadByEntitiesIfPresent(
                           ImmutableList.of(theEntity, new TestEntity("does-not-exist", "bar")));
               assertThat(results).containsExactly(theEntity);
-              assertDetached(results.get(0));
+              assertDetachedFromEntityManager(results.get(0));
             });
   }
 
@@ -491,7 +499,7 @@ class JpaTransactionManagerImplTest {
               ImmutableList<TestEntity> results =
                   jpaTm().loadByEntities(ImmutableList.of(theEntity));
               assertThat(results).containsExactly(theEntity);
-              assertDetached(results.get(0));
+              assertDetachedFromEntityManager(results.get(0));
             });
   }
 
@@ -503,7 +511,7 @@ class JpaTransactionManagerImplTest {
             .transact(
                 () ->
                     jpaTm().loadAllOf(TestEntity.class).stream()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(DatabaseHelper::assertDetachedFromEntityManager)
                         .collect(toImmutableList()));
     assertThat(persisted).containsExactlyElementsIn(moreEntities);
   }

core/src/test/java/google/registry/persistence/transaction/QueryComposerTest.java

@@ -77,7 +77,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GT, "bravo")
                         .first()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .get()))
         .isEqualTo(charlie);
     assertThat(
@@ -86,7 +86,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GTE, "charlie")
                         .first()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .get()))
         .isEqualTo(charlie);
     assertThat(
@@ -95,7 +95,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LT, "bravo")
                         .first()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .get()))
         .isEqualTo(alpha);
     assertThat(
@@ -104,7 +104,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LTE, "alpha")
                         .first()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .get()))
         .isEqualTo(alpha);
   }
@@ -125,7 +125,7 @@ public class QueryComposerTest {
     assertThat(
             transactIfJpaTm(
                 () ->
-                    DatabaseHelper.assertDetached(
+                    QueryComposerTest.assertDetachedIfJpa(
                         tm().createQueryComposer(TestEntity.class)
                             .where("name", Comparator.EQ, "alpha")
                             .getSingleResult())))
@@ -174,7 +174,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GT, "alpha")
                         .stream()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .collect(toImmutableList())))
         .containsExactly(bravo, charlie);
     assertThat(
@@ -184,7 +184,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.GTE, "bravo")
                         .stream()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .collect(toImmutableList())))
         .containsExactly(bravo, charlie);
     assertThat(
@@ -194,7 +194,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LT, "charlie")
                         .stream()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .collect(toImmutableList())))
         .containsExactly(alpha, bravo);
     assertThat(
@@ -204,7 +204,7 @@ public class QueryComposerTest {
                         .createQueryComposer(TestEntity.class)
                         .where("name", Comparator.LTE, "bravo")
                         .stream()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .collect(toImmutableList())))
         .containsExactly(alpha, bravo);
   }
@@ -228,7 +228,7 @@ public class QueryComposerTest {
                     tm().createQueryComposer(TestEntity.class)
                         .where("val", Comparator.EQ, 2)
                         .first()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .get()))
         .isEqualTo(bravo);
   }
@@ -243,7 +243,7 @@ public class QueryComposerTest {
                         .where("val", Comparator.GT, 1)
                         .orderBy("val")
                         .stream()
-                        .map(DatabaseHelper::assertDetached)
+                        .map(QueryComposerTest::assertDetachedIfJpa)
                         .collect(toImmutableList())))
         .containsExactly(bravo, alpha);
   }
@@ -354,6 +354,13 @@ public class QueryComposerTest {
         .contains("The LIKE operation is not supported on Datastore.");
   }
 
+  private static <T> T assertDetachedIfJpa(T entity) {
+    if (!tm().isOfy()) {
+      return DatabaseHelper.assertDetachedFromEntityManager(entity);
+    }
+    return entity;
+  }
+
   @javax.persistence.Entity
   @Entity(name = "QueryComposerTestEntity")
   private static class TestEntity extends ImmutableObject {

core/src/test/java/google/registry/testing/DatabaseHelper.java

@@ -1351,14 +1351,12 @@ public class DatabaseHelper {
   }
 
   /**
-   * In JPA mode, assert that the given entity is detached from the current entity manager.
+   * Asserts that the given entity is detached from the current JPA entity manager.
    *
    * <p>Returns the original entity object.
    */
-  public static <T> T assertDetached(T entity) {
-    if (!tm().isOfy()) {
-      assertThat(jpaTm().getEntityManager().contains(entity)).isFalse();
-    }
+  public static <T> T assertDetachedFromEntityManager(T entity) {
+    assertThat(jpaTm().getEntityManager().contains(entity)).isFalse();
     return entity;
   }