Add JpaTransactionManagerRule (#277)

* Add RegistryRuntime and JpaTransactionManagerRule

* Revert RegistryJavaRuntime change

* Add JpaTransactionManager interface

core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java

@@ -14,7 +14,6 @@
 
 package google.registry.model.ofy;
 
-
 import static google.registry.model.ofy.ObjectifyService.ofy;
 
 import google.registry.model.transaction.TransactionManager;

core/src/main/java/google/registry/model/transaction/DummyJpaTransactionManager.java

@@ -0,0 +1,44 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.transaction;
+
+import java.lang.reflect.Proxy;
+
+/**
+ * A dummy implementation for {@link JpaTransactionManager} which throws exception when any of its
+ * method is invoked.
+ *
+ * <p>This is used to initialize the {@link TransactionManagerFactory#jpaTm} when running unit
+ * tests, because obviously we cannot connect to the actual Cloud SQL backend in a unit test.
+ *
+ * <p>If a unit test needs to access the Cloud SQL database, it must add JpaTransactionManagerRule
+ * as a JUnit rule in the test class.
+ */
+public class DummyJpaTransactionManager {
+
+  /** Constructs a dummy {@link JpaTransactionManager} instance. */
+  public static JpaTransactionManager create() {
+    return (JpaTransactionManager)
+        Proxy.newProxyInstance(
+            JpaTransactionManager.class.getClassLoader(),
+            new Class[] {JpaTransactionManager.class},
+            (proxy, method, args) -> {
+              throw new UnsupportedOperationException(
+                  "JpaTransactionManager was not initialized as the runtime is detected as"
+                      + " Unittest. Add JpaTransactionManagerRule in the unit test for"
+                      + " initialization.");
+            });
+  }
+}

core/src/main/java/google/registry/model/transaction/JpaTransactionManager.java

@@ -14,154 +14,10 @@
 
 package google.registry.model.transaction;
 
-import com.google.common.flogger.FluentLogger;
-import google.registry.persistence.PersistenceModule.AppEngineEmf;
-import google.registry.util.Clock;
-import javax.inject.Inject;
-import javax.inject.Singleton;
 import javax.persistence.EntityManager;
-import javax.persistence.EntityManagerFactory;
-import javax.persistence.EntityTransaction;
-import javax.persistence.PersistenceException;
-import org.joda.time.DateTime;
 
-/** Implementation of {@link TransactionManager} for JPA compatible database. */
+/** Sub-interface of {@link TransactionManager} which defines JPA related methods. */
-@Singleton
+public interface JpaTransactionManager extends TransactionManager {
-public class JpaTransactionManager implements TransactionManager {
+  /** Returns the {@link EntityManager} for the current request. */
-
+  EntityManager getEntityManager();
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
-  // EntityManagerFactory is thread safe.
-  private final EntityManagerFactory emf;
-  private final Clock clock;
-  // TODO(shicong): Investigate alternatives for managing transaction information. ThreadLocal adds
-  //  an unnecessary restriction that each request has to be processed by one thread synchronously.
-  private final ThreadLocal<TransactionInfo> transactionInfo =
-      ThreadLocal.withInitial(TransactionInfo::new);
-
-  @Inject
-  JpaTransactionManager(@AppEngineEmf EntityManagerFactory emf, Clock clock) {
-    this.emf = emf;
-    this.clock = clock;
-  }
-
-  public EntityManager getEntityManager() {
-    if (transactionInfo.get().entityManager == null) {
-      throw new PersistenceException(
-          "No EntityManager has been initialized. getEntityManager() must be invoked in the scope"
-              + " of a transaction");
-    }
-    return transactionInfo.get().entityManager;
-  }
-
-  @Override
-  public boolean inTransaction() {
-    return transactionInfo.get().inTransaction;
-  }
-
-  @Override
-  public void assertInTransaction() {
-    if (!inTransaction()) {
-      throw new PersistenceException("Not in a transaction");
-    }
-  }
-
-  @Override
-  public <T> T transact(Work<T> work) {
-    // TODO(shicong): Investigate removing transactNew functionality after migration as it may
-    //  be same as this one.
-    if (inTransaction()) {
-      return work.run();
-    }
-    TransactionInfo txnInfo = transactionInfo.get();
-    txnInfo.entityManager = emf.createEntityManager();
-    EntityTransaction txn = txnInfo.entityManager.getTransaction();
-    try {
-      txn.begin();
-      txnInfo.inTransaction = true;
-      txnInfo.transactionTime = clock.nowUtc();
-      T result = work.run();
-      txn.commit();
-      return result;
-    } catch (Throwable transactionException) {
-      String rollbackMessage;
-      try {
-        txn.rollback();
-        rollbackMessage = "transaction rolled back";
-      } catch (Throwable rollbackException) {
-        logger.atSevere().withCause(rollbackException).log("Rollback failed, suppressing error");
-        rollbackMessage = "transaction rollback failed";
-      }
-      throw new PersistenceException(
-          "Error during transaction, " + rollbackMessage, transactionException);
-    } finally {
-      txnInfo.clear();
-    }
-  }
-
-  @Override
-  public void transact(Runnable work) {
-    transact(
-        () -> {
-          work.run();
-          return null;
-        });
-  }
-
-  @Override
-  public <T> T transactNew(Work<T> work) {
-    // TODO(shicong): Implements the functionality to start a new transaction.
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
-  public void transactNew(Runnable work) {
-    // TODO(shicong): Implements the functionality to start a new transaction.
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
-  public <T> T transactNewReadOnly(Work<T> work) {
-    // TODO(shicong): Implements read only transaction.
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
-  public void transactNewReadOnly(Runnable work) {
-    // TODO(shicong): Implements read only transaction.
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
-  public <T> T doTransactionless(Work<T> work) {
-    // TODO(shicong): Implements doTransactionless.
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
-  public DateTime getTransactionTime() {
-    assertInTransaction();
-    TransactionInfo txnInfo = transactionInfo.get();
-    if (txnInfo.transactionTime == null) {
-      throw new PersistenceException("In a transaction but transactionTime is null");
-    }
-    return txnInfo.transactionTime;
-  }
-
-  private static class TransactionInfo {
-    EntityManager entityManager;
-    boolean inTransaction = false;
-    DateTime transactionTime;
-
-    private void clear() {
-      inTransaction = false;
-      transactionTime = null;
-      if (entityManager != null) {
-        // Close this EntityManager just let the connection pool be able to reuse it, it doesn't
-        // close the underlying database connection.
-        entityManager.close();
-        entityManager = null;
-      }
-    }
-  }
 }

core/src/main/java/google/registry/model/transaction/JpaTransactionManagerImpl.java

@@ -0,0 +1,168 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.transaction;
+
+import com.google.common.flogger.FluentLogger;
+import google.registry.persistence.PersistenceModule.AppEngineEmf;
+import google.registry.util.Clock;
+import javax.inject.Inject;
+import javax.inject.Singleton;
+import javax.persistence.EntityManager;
+import javax.persistence.EntityManagerFactory;
+import javax.persistence.EntityTransaction;
+import javax.persistence.PersistenceException;
+import org.joda.time.DateTime;
+
+/** Implementation of {@link JpaTransactionManager} for JPA compatible database. */
+@Singleton
+public class JpaTransactionManagerImpl implements JpaTransactionManager {
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  // EntityManagerFactory is thread safe.
+  private final EntityManagerFactory emf;
+  private final Clock clock;
+  // TODO(shicong): Investigate alternatives for managing transaction information. ThreadLocal adds
+  //  an unnecessary restriction that each request has to be processed by one thread synchronously.
+  private final ThreadLocal<TransactionInfo> transactionInfo =
+      ThreadLocal.withInitial(TransactionInfo::new);
+
+  @Inject
+  JpaTransactionManagerImpl(@AppEngineEmf EntityManagerFactory emf, Clock clock) {
+    this.emf = emf;
+    this.clock = clock;
+  }
+
+  @Override
+  public EntityManager getEntityManager() {
+    if (transactionInfo.get().entityManager == null) {
+      throw new PersistenceException(
+          "No EntityManager has been initialized. getEntityManager() must be invoked in the scope"
+              + " of a transaction");
+    }
+    return transactionInfo.get().entityManager;
+  }
+
+  @Override
+  public boolean inTransaction() {
+    return transactionInfo.get().inTransaction;
+  }
+
+  @Override
+  public void assertInTransaction() {
+    if (!inTransaction()) {
+      throw new PersistenceException("Not in a transaction");
+    }
+  }
+
+  @Override
+  public <T> T transact(Work<T> work) {
+    // TODO(shicong): Investigate removing transactNew functionality after migration as it may
+    //  be same as this one.
+    if (inTransaction()) {
+      return work.run();
+    }
+    TransactionInfo txnInfo = transactionInfo.get();
+    txnInfo.entityManager = emf.createEntityManager();
+    EntityTransaction txn = txnInfo.entityManager.getTransaction();
+    try {
+      txn.begin();
+      txnInfo.inTransaction = true;
+      txnInfo.transactionTime = clock.nowUtc();
+      T result = work.run();
+      txn.commit();
+      return result;
+    } catch (Throwable transactionException) {
+      String rollbackMessage;
+      try {
+        txn.rollback();
+        rollbackMessage = "transaction rolled back";
+      } catch (Throwable rollbackException) {
+        logger.atSevere().withCause(rollbackException).log("Rollback failed, suppressing error");
+        rollbackMessage = "transaction rollback failed";
+      }
+      throw new PersistenceException(
+          "Error during transaction, " + rollbackMessage, transactionException);
+    } finally {
+      txnInfo.clear();
+    }
+  }
+
+  @Override
+  public void transact(Runnable work) {
+    transact(
+        () -> {
+          work.run();
+          return null;
+        });
+  }
+
+  @Override
+  public <T> T transactNew(Work<T> work) {
+    // TODO(shicong): Implements the functionality to start a new transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void transactNew(Runnable work) {
+    // TODO(shicong): Implements the functionality to start a new transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public <T> T transactNewReadOnly(Work<T> work) {
+    // TODO(shicong): Implements read only transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void transactNewReadOnly(Runnable work) {
+    // TODO(shicong): Implements read only transaction.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public <T> T doTransactionless(Work<T> work) {
+    // TODO(shicong): Implements doTransactionless.
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public DateTime getTransactionTime() {
+    assertInTransaction();
+    TransactionInfo txnInfo = transactionInfo.get();
+    if (txnInfo.transactionTime == null) {
+      throw new PersistenceException("In a transaction but transactionTime is null");
+    }
+    return txnInfo.transactionTime;
+  }
+
+  private static class TransactionInfo {
+    EntityManager entityManager;
+    boolean inTransaction = false;
+    DateTime transactionTime;
+
+    private void clear() {
+      inTransaction = false;
+      transactionTime = null;
+      if (entityManager != null) {
+        // Close this EntityManager just let the connection pool be able to reuse it, it doesn't
+        // close the underlying database connection.
+        entityManager.close();
+        entityManager = null;
+      }
+    }
+  }
+}

core/src/main/java/google/registry/model/transaction/TransactionManagerFactory.java

@@ -14,21 +14,29 @@
 
 package google.registry.model.transaction;
 
+import com.google.appengine.api.utils.SystemProperty;
+import com.google.appengine.api.utils.SystemProperty.Environment.Value;
 import com.google.common.annotations.VisibleForTesting;
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.DaggerPersistenceComponent;
-import google.registry.persistence.PersistenceComponent;
 
 /** Factory class to create {@link TransactionManager} instance. */
 // TODO: Rename this to PersistenceFactory and move to persistence package.
 public class TransactionManagerFactory {
 
   private static final TransactionManager TM = createTransactionManager();
-
+  @VisibleForTesting static JpaTransactionManager jpaTm = createJpaTransactionManager();
-  @VisibleForTesting static PersistenceComponent component = DaggerPersistenceComponent.create();
 
   private TransactionManagerFactory() {}
 
+  private static JpaTransactionManager createJpaTransactionManager() {
+    if (SystemProperty.environment.value() == Value.Production) {
+      return DaggerPersistenceComponent.create().jpaTransactionManager();
+    } else {
+      return DummyJpaTransactionManager.create();
+    }
+  }
+
   private static TransactionManager createTransactionManager() {
     // TODO: Determine how to provision TransactionManager after the dual-write. During the
     // dual-write transitional phase, we need the TransactionManager for both Datastore and Cloud
@@ -48,6 +56,6 @@ public class TransactionManagerFactory {
     //    1. App Engine
     //    2. Local JVM used by nomulus tool
     //    3. Unit test
-    return component.jpaTransactionManager();
+    return jpaTm;
   }
 }

core/src/main/java/google/registry/persistence/PersistenceComponent.java

@@ -18,7 +18,7 @@ import dagger.Component;
 import google.registry.config.CredentialModule;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.keyring.kms.KmsModule;
-import google.registry.model.transaction.JpaTransactionManager;
+import google.registry.model.transaction.JpaTransactionManagerImpl;
 import google.registry.util.UtilsModule;
 import javax.inject.Singleton;
 import javax.persistence.EntityManagerFactory;
@@ -34,5 +34,5 @@ import javax.persistence.EntityManagerFactory;
       UtilsModule.class
     })
 public interface PersistenceComponent {
-  JpaTransactionManager jpaTransactionManager();
+  JpaTransactionManagerImpl jpaTransactionManager();
 }

core/src/test/java/google/registry/model/transaction/DummyJpaTransactionManagerTest.java

@@ -0,0 +1,33 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.transaction;
+
+import static google.registry.model.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.JUnitBackports.assertThrows;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** JUnit test for {@link DummyJpaTransactionManager} */
+@RunWith(JUnit4.class)
+public class DummyJpaTransactionManagerTest {
+
+  @Test
+  public void throwsExceptionWhenAnyMethodIsInvoked() {
+    assertThrows(UnsupportedOperationException.class, () -> jpaTm().transact(() -> null));
+    assertThrows(UnsupportedOperationException.class, () -> jpaTm().getTransactionTime());
+  }
+}

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

@@ -0,0 +1,192 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.transaction;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.JUnitBackports.assertThrows;
+import static google.registry.testing.TestDataHelper.fileClassPath;
+
+import google.registry.testing.FakeClock;
+import java.math.BigInteger;
+import javax.persistence.EntityManager;
+import javax.persistence.PersistenceException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link JpaTransactionManagerImpl}. */
+@RunWith(JUnit4.class)
+public class JpaTransactionManagerImplTest {
+  @Rule
+  public final JpaTransactionManagerRule jpaTmRule =
+      new JpaTransactionManagerRule.Builder()
+          .withInitScript(fileClassPath(getClass(), "test_schema.sql"))
+          .build();
+
+  @Test
+  public void inTransaction_returnsCorrespondingResult() {
+    assertThat(jpaTm().inTransaction()).isFalse();
+    jpaTm().transact(() -> assertThat(jpaTm().inTransaction()).isTrue());
+    assertThat(jpaTm().inTransaction()).isFalse();
+  }
+
+  @Test
+  public void assertInTransaction_throwsExceptionWhenNotInTransaction() {
+    assertThrows(PersistenceException.class, () -> jpaTm().assertInTransaction());
+    jpaTm().transact(() -> jpaTm().assertInTransaction());
+    assertThrows(PersistenceException.class, () -> jpaTm().assertInTransaction());
+  }
+
+  @Test
+  public void getTransactionTime_throwsExceptionWhenNotInTransaction() {
+    FakeClock txnClock = jpaTmRule.getTxnClock();
+    txnClock.advanceOneMilli();
+    assertThrows(PersistenceException.class, () -> jpaTm().getTransactionTime());
+    jpaTm().transact(() -> assertThat(jpaTm().getTransactionTime()).isEqualTo(txnClock.nowUtc()));
+    assertThrows(PersistenceException.class, () -> jpaTm().getTransactionTime());
+  }
+
+  @Test
+  public void transact_succeeds() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    jpaTm()
+        .transact(
+            () -> {
+              insertPerson(10);
+              insertCompany("Foo");
+              insertCompany("Bar");
+            });
+    assertPersonCount(1);
+    assertPersonExist(10);
+    assertCompanyCount(2);
+    assertCompanyExist("Foo");
+    assertCompanyExist("Bar");
+  }
+
+  @Test
+  public void transact_hasNoEffectWithPartialSuccess() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    assertThrows(
+        RuntimeException.class,
+        () ->
+            jpaTm()
+                .transact(
+                    () -> {
+                      insertPerson(10);
+                      insertCompany("Foo");
+                      throw new RuntimeException();
+                    }));
+    assertPersonEmpty();
+    assertCompanyEmpty();
+  }
+
+  @Test
+  public void transact_reusesExistingTransaction() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    jpaTm()
+        .transact(
+            () ->
+                jpaTm()
+                    .transact(
+                        () -> {
+                          insertPerson(10);
+                          insertCompany("Foo");
+                          insertCompany("Bar");
+                        }));
+    assertPersonCount(1);
+    assertPersonExist(10);
+    assertCompanyCount(2);
+    assertCompanyExist("Foo");
+    assertCompanyExist("Bar");
+  }
+
+  private void insertPerson(int age) {
+    jpaTm()
+        .getEntityManager()
+        .createNativeQuery(String.format("INSERT INTO Person (age) VALUES (%d)", age))
+        .executeUpdate();
+  }
+
+  private void insertCompany(String name) {
+    jpaTm()
+        .getEntityManager()
+        .createNativeQuery(String.format("INSERT INTO Company (name) VALUES ('%s')", name))
+        .executeUpdate();
+  }
+
+  private void assertPersonExist(int age) {
+    jpaTm()
+        .transact(
+            () -> {
+              EntityManager em = jpaTm().getEntityManager();
+              Integer maybeAge =
+                  (Integer)
+                      em.createNativeQuery(
+                              String.format("SELECT age FROM Person WHERE age = %d", age))
+                          .getSingleResult();
+              assertThat(maybeAge).isEqualTo(age);
+            });
+  }
+
+  private void assertCompanyExist(String name) {
+    jpaTm()
+        .transact(
+            () -> {
+              String maybeName =
+                  (String)
+                      jpaTm()
+                          .getEntityManager()
+                          .createNativeQuery(
+                              String.format("SELECT name FROM Company WHERE name = '%s'", name))
+                          .getSingleResult();
+              assertThat(maybeName).isEqualTo(name);
+            });
+  }
+
+  private void assertPersonCount(int count) {
+    assertThat(countTable("Person")).isEqualTo(count);
+  }
+
+  private void assertCompanyCount(int count) {
+    assertThat(countTable("Company")).isEqualTo(count);
+  }
+
+  private void assertPersonEmpty() {
+    assertPersonCount(0);
+  }
+
+  private void assertCompanyEmpty() {
+    assertCompanyCount(0);
+  }
+
+  private int countTable(String tableName) {
+    return jpaTm()
+        .transact(
+            () -> {
+              BigInteger colCount =
+                  (BigInteger)
+                      jpaTm()
+                          .getEntityManager()
+                          .createNativeQuery(String.format("SELECT COUNT(*) FROM %s", tableName))
+                          .getSingleResult();
+              return colCount.intValue();
+            });
+  }
+}

core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java

@@ -0,0 +1,109 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.transaction;
+
+import static org.joda.time.DateTimeZone.UTC;
+
+import google.registry.persistence.PersistenceModule;
+import google.registry.testing.FakeClock;
+import javax.persistence.EntityManagerFactory;
+import org.joda.time.DateTime;
+import org.junit.rules.ExternalResource;
+import org.junit.rules.RuleChain;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+import org.testcontainers.containers.JdbcDatabaseContainer;
+import org.testcontainers.containers.PostgreSQLContainer;
+
+/**
+ * JUnit Rule to provision {@link JpaTransactionManagerImpl} backed by {@link PostgreSQLContainer}.
+ *
+ * <p>This rule also replaces the {@link JpaTransactionManagerImpl} provided by {@link
+ * TransactionManagerFactory} with the {@link JpaTransactionManagerImpl} generated by the rule
+ * itself, so that all SQL queries will be sent to the database instance created by {@link
+ * PostgreSQLContainer} to achieve test purpose.
+ */
+public class JpaTransactionManagerRule extends ExternalResource {
+  private static final String SCHEMA_GOLDEN_SQL = "sql/schema/nomulus.golden.sql";
+
+  private final DateTime now = DateTime.now(UTC);
+  private final FakeClock clock = new FakeClock(now);
+  private final String initScript;
+  private JdbcDatabaseContainer database;
+  private EntityManagerFactory emf;
+  private JpaTransactionManager cachedTm;
+
+  private JpaTransactionManagerRule(String initScript) {
+    this.initScript = initScript;
+  }
+
+  /** Wraps {@link JpaTransactionManagerRule} in a {@link PostgreSQLContainer}. */
+  @Override
+  public Statement apply(Statement base, Description description) {
+    database = new PostgreSQLContainer().withInitScript(initScript);
+    return RuleChain.outerRule(database)
+        .around(JpaTransactionManagerRule.super::apply)
+        .apply(base, description);
+  }
+
+  @Override
+  public void before() {
+    emf =
+        PersistenceModule.create(
+            database.getJdbcUrl(),
+            database.getUsername(),
+            database.getPassword(),
+            PersistenceModule.providesDefaultDatabaseConfigs());
+    JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock);
+    cachedTm = TransactionManagerFactory.jpaTm;
+    TransactionManagerFactory.jpaTm = txnManager;
+  }
+
+  @Override
+  public void after() {
+    TransactionManagerFactory.jpaTm = cachedTm;
+    if (emf != null) {
+      emf.close();
+    }
+    cachedTm = null;
+  }
+
+  /** Returns the {@link FakeClock} used by the underlying {@link JpaTransactionManagerImpl}. */
+  public FakeClock getTxnClock() {
+    return clock;
+  }
+
+  /** Builder for {@link JpaTransactionManagerRule}. */
+  public static class Builder {
+    private String initScript;
+
+    /**
+     * Sets the SQL script to be used to initialize the database. If not set,
+     * sql/schema/nomulus.golden.sql will be used.
+     */
+    public Builder withInitScript(String initScript) {
+      this.initScript = initScript;
+      return this;
+    }
+
+    /** Builds a {@link JpaTransactionManagerRule} instance. */
+    public JpaTransactionManagerRule build() {
+      if (initScript == null) {
+        initScript = SCHEMA_GOLDEN_SQL;
+      }
+      return new JpaTransactionManagerRule(initScript);
+    }
+  }
+}

core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRuleTest.java

@@ -0,0 +1,59 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.transaction;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.testing.JUnitBackports.assertThrows;
+
+import java.util.List;
+import javax.persistence.PersistenceException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** JUnit test for {@link JpaTransactionManagerRule} */
+@RunWith(JUnit4.class)
+public class JpaTransactionManagerRuleTest {
+
+  @Rule
+  public final JpaTransactionManagerRule jpaTmRule =
+      new JpaTransactionManagerRule.Builder().build();
+
+  @Test
+  public void verifiesRuleWorks() {
+    assertThrows(
+        PersistenceException.class,
+        () ->
+            jpaTm()
+                .transact(
+                    () ->
+                        jpaTm()
+                            .getEntityManager()
+                            .createNativeQuery("SELECT * FROM NoneExistentTable")
+                            .getResultList()));
+    jpaTm()
+        .transact(
+            () -> {
+              List results =
+                  jpaTm()
+                      .getEntityManager()
+                      .createNativeQuery("SELECT * FROM \"ClaimsList\"")
+                      .getResultList();
+              assertThat(results).isEmpty();
+            });
+  }
+}

core/src/test/java/google/registry/model/transaction/JpaTransactionManagerTest.java

@@ -1,211 +0,0 @@
-// Copyright 2019 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.model.transaction;
-
-import static com.google.common.truth.Truth.assertThat;
-import static google.registry.testing.JUnitBackports.assertThrows;
-import static google.registry.testing.TestDataHelper.fileClassPath;
-import static org.joda.time.DateTimeZone.UTC;
-
-import google.registry.persistence.PersistenceModule;
-import google.registry.testing.FakeClock;
-import google.registry.util.Clock;
-import java.math.BigInteger;
-import javax.persistence.EntityManager;
-import javax.persistence.EntityManagerFactory;
-import javax.persistence.PersistenceException;
-import org.joda.time.DateTime;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.testcontainers.containers.JdbcDatabaseContainer;
-import org.testcontainers.containers.PostgreSQLContainer;
-
-/** Unit tests for {@link JpaTransactionManager}. */
-@RunWith(JUnit4.class)
-public class JpaTransactionManagerTest {
-  @Rule
-  public JdbcDatabaseContainer database =
-      new PostgreSQLContainer().withInitScript(fileClassPath(getClass(), "test_schema.sql"));
-
-  private DateTime now = DateTime.now(UTC);
-  private Clock clock = new FakeClock(now);
-  private EntityManagerFactory emf;
-  private JpaTransactionManager txnManager;
-
-  @Before
-  public void init() {
-    emf =
-        PersistenceModule.create(
-            database.getJdbcUrl(),
-            database.getUsername(),
-            database.getPassword(),
-            PersistenceModule.providesDefaultDatabaseConfigs());
-    txnManager = new JpaTransactionManager(emf, clock);
-  }
-
-  @After
-  public void clear() {
-    if (emf != null) {
-      emf.close();
-    }
-  }
-
-  @Test
-  public void inTransaction_returnsCorrespondingResult() {
-    assertThat(txnManager.inTransaction()).isFalse();
-    txnManager.transact(() -> assertThat(txnManager.inTransaction()).isTrue());
-    assertThat(txnManager.inTransaction()).isFalse();
-  }
-
-  @Test
-  public void assertInTransaction_throwsExceptionWhenNotInTransaction() {
-    assertThrows(PersistenceException.class, () -> txnManager.assertInTransaction());
-    txnManager.transact(() -> txnManager.assertInTransaction());
-    assertThrows(PersistenceException.class, () -> txnManager.assertInTransaction());
-  }
-
-  @Test
-  public void getTransactionTime_throwsExceptionWhenNotInTransaction() {
-    assertThrows(PersistenceException.class, () -> txnManager.getTransactionTime());
-    txnManager.transact(() -> assertThat(txnManager.getTransactionTime()).isEqualTo(now));
-    assertThrows(PersistenceException.class, () -> txnManager.getTransactionTime());
-  }
-
-  @Test
-  public void transact_succeeds() {
-    assertPersonEmpty();
-    assertCompanyEmpty();
-    txnManager.transact(
-        () -> {
-          insertPerson(10);
-          insertCompany("Foo");
-          insertCompany("Bar");
-        });
-    assertPersonCount(1);
-    assertPersonExist(10);
-    assertCompanyCount(2);
-    assertCompanyExist("Foo");
-    assertCompanyExist("Bar");
-  }
-
-  @Test
-  public void transact_hasNoEffectWithPartialSuccess() {
-    assertPersonEmpty();
-    assertCompanyEmpty();
-    assertThrows(
-        RuntimeException.class,
-        () ->
-            txnManager.transact(
-                () -> {
-                  insertPerson(10);
-                  insertCompany("Foo");
-                  throw new RuntimeException();
-                }));
-    assertPersonEmpty();
-    assertCompanyEmpty();
-  }
-
-  @Test
-  public void transact_reusesExistingTransaction() {
-    assertPersonEmpty();
-    assertCompanyEmpty();
-    txnManager.transact(
-        () ->
-            txnManager.transact(
-                () -> {
-                  insertPerson(10);
-                  insertCompany("Foo");
-                  insertCompany("Bar");
-                }));
-    assertPersonCount(1);
-    assertPersonExist(10);
-    assertCompanyCount(2);
-    assertCompanyExist("Foo");
-    assertCompanyExist("Bar");
-  }
-
-  private void insertPerson(int age) {
-    txnManager
-        .getEntityManager()
-        .createNativeQuery(String.format("INSERT INTO Person (age) VALUES (%d)", age))
-        .executeUpdate();
-  }
-
-  private void insertCompany(String name) {
-    txnManager
-        .getEntityManager()
-        .createNativeQuery(String.format("INSERT INTO Company (name) VALUES ('%s')", name))
-        .executeUpdate();
-  }
-
-  private void assertPersonExist(int age) {
-    txnManager.transact(
-        () -> {
-          EntityManager em = txnManager.getEntityManager();
-          Integer maybeAge =
-              (Integer)
-                  em.createNativeQuery(String.format("SELECT age FROM Person WHERE age = %d", age))
-                      .getSingleResult();
-          assertThat(maybeAge).isEqualTo(age);
-        });
-  }
-
-  private void assertCompanyExist(String name) {
-    txnManager.transact(
-        () -> {
-          String maybeName =
-              (String)
-                  txnManager
-                      .getEntityManager()
-                      .createNativeQuery(
-                          String.format("SELECT name FROM Company WHERE name = '%s'", name))
-                      .getSingleResult();
-          assertThat(maybeName).isEqualTo(name);
-        });
-  }
-
-  private void assertPersonCount(int count) {
-    assertThat(countTable("Person")).isEqualTo(count);
-  }
-
-  private void assertCompanyCount(int count) {
-    assertThat(countTable("Company")).isEqualTo(count);
-  }
-
-  private void assertPersonEmpty() {
-    assertPersonCount(0);
-  }
-
-  private void assertCompanyEmpty() {
-    assertCompanyCount(0);
-  }
-
-  private int countTable(String tableName) {
-    return txnManager.transact(
-        () -> {
-          BigInteger colCount =
-              (BigInteger)
-                  txnManager
-                      .getEntityManager()
-                      .createNativeQuery(String.format("SELECT COUNT(*) FROM %s", tableName))
-                      .getSingleResult();
-          return colCount.intValue();
-        });
-  }
-}