Convert 3 classes from ofy -> tm (#1034)

* Convert 3 classes from ofy -> tm

Convert SyncGroupMembersAction, SyncRegistrarsSheet and
IcannReportingUploadAction and their test cases to use TransactionManager and
dual-test them so we know they work in jpa.

* Address comments in review

Address review comments and make the entire IcannReportingUploadAction run
transactional.

* reformatted.

* Remove duplicate loadByKey() method

Remove test method added in a recent PR.

core/src/main/java/google/registry/export/SyncGroupMembersAction.java

@@ -16,7 +16,6 @@ package google.registry.export;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.util.CollectionUtils.nullToEmpty;
@@ -164,7 +163,7 @@ public final class SyncGroupMembersAction implements Runnable {
         registrarsToSave.add(result.getKey().asBuilder().setContactsRequireSyncing(false).build());
       }
     }
-    tm().transactNew(() -> ofy().save().entities(registrarsToSave.build()));
+    tm().transactNew(() -> tm().updateAll(registrarsToSave.build()));
     return errors;
   }
 

core/src/main/java/google/registry/export/sheet/SyncRegistrarsSheet.java

@@ -17,7 +17,6 @@ package google.registry.export.sheet;
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.common.Cursor.CursorType.SYNC_REGISTRAR_SHEET;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registrar.RegistrarContact.Type.ABUSE;
 import static google.registry.model.registrar.RegistrarContact.Type.ADMIN;
 import static google.registry.model.registrar.RegistrarContact.Type.BILLING;
@@ -26,6 +25,7 @@ import static google.registry.model.registrar.RegistrarContact.Type.MARKETING;
 import static google.registry.model.registrar.RegistrarContact.Type.TECH;
 import static google.registry.model.registrar.RegistrarContact.Type.WHOIS;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.google.common.base.Joiner;
@@ -40,6 +40,7 @@ import google.registry.model.registrar.RegistrarContact;
 import google.registry.util.Clock;
 import google.registry.util.DateTimeUtils;
 import java.io.IOException;
+import java.util.Optional;
 import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
@@ -61,8 +62,10 @@ class SyncRegistrarsSheet {
    * successfully completed, as measured by a cursor.
    */
   boolean wereRegistrarsModified() {
-    Cursor cursor = ofy().load().key(Cursor.createGlobalKey(SYNC_REGISTRAR_SHEET)).now();
+    Optional<Cursor> cursor =
-    DateTime lastUpdateTime = (cursor == null) ? START_OF_TIME : cursor.getCursorTime();
+        transactIfJpaTm(
+            () -> tm().loadByKeyIfPresent(Cursor.createGlobalVKey(SYNC_REGISTRAR_SHEET)));
+    DateTime lastUpdateTime = !cursor.isPresent() ? START_OF_TIME : cursor.get().getCursorTime();
     for (Registrar registrar : Registrar.loadAllCached()) {
       if (DateTimeUtils.isAtOrAfter(registrar.getLastUpdateTime(), lastUpdateTime)) {
         return true;

core/src/main/java/google/registry/reporting/icann/IcannReportingUploadAction.java

@@ -17,7 +17,6 @@ package google.registry.reporting.icann;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
@@ -28,7 +27,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.io.ByteStreams;
-import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.gcs.GcsUtils;
 import google.registry.model.common.Cursor;
@@ -36,6 +34,7 @@ import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.registry.Registries;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
+import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.HttpException.ServiceUnavailableException;
 import google.registry.request.Response;
@@ -100,7 +99,7 @@ public final class IcannReportingUploadAction implements Runnable {
 
   @Override
   public void run() {
-    Callable<Void> lockRunner =
+    Runnable transactional =
         () -> {
           ImmutableMap.Builder<String, Boolean> reportSummaryBuilder = new ImmutableMap.Builder<>();
 
@@ -122,6 +121,11 @@ public final class IcannReportingUploadAction implements Runnable {
           emailUploadResults(reportSummaryBuilder.build());
           response.setStatus(SC_OK);
           response.setContentType(PLAIN_TEXT_UTF_8);
+        };
+
+    Callable<Void> lockRunner =
+        () -> {
+          tm().transact(transactional);
           return null;
         };
 
@@ -185,7 +189,7 @@ public final class IcannReportingUploadAction implements Runnable {
               cursorType,
               cursorTime.withTimeAtStartOfDay().withDayOfMonth(1).plusMonths(1),
               Registry.get(tldStr));
-      tm().transact(() -> tm().put(newCursor));
+      tm().put(newCursor);
     }
   }
 
@@ -204,16 +208,16 @@ public final class IcannReportingUploadAction implements Runnable {
 
     ImmutableSet<Registry> registries = Registries.getTldEntitiesOfType(TldType.REAL);
 
-    ImmutableMap<Key<Cursor>, Registry> activityKeyMap =
+    ImmutableMap<VKey<? extends Cursor>, Registry> activityKeyMap =
         loadKeyMap(registries, CursorType.ICANN_UPLOAD_ACTIVITY);
-    ImmutableMap<Key<Cursor>, Registry> transactionKeyMap =
+    ImmutableMap<VKey<? extends Cursor>, Registry> transactionKeyMap =
         loadKeyMap(registries, CursorType.ICANN_UPLOAD_TX);
 
-    ImmutableSet.Builder<Key<Cursor>> keys = new ImmutableSet.Builder<>();
+    ImmutableSet.Builder<VKey<? extends Cursor>> keys = new ImmutableSet.Builder<>();
     keys.addAll(activityKeyMap.keySet());
     keys.addAll(transactionKeyMap.keySet());
 
-    Map<Key<Cursor>, Cursor> cursorMap = ofy().load().keys(keys.build());
+    Map<VKey<? extends Cursor>, Cursor> cursorMap = tm().loadByKeysIfPresent(keys.build());
     ImmutableMap.Builder<Cursor, String> cursors = new ImmutableMap.Builder<>();
     cursors.putAll(
         defaultNullCursorsToNextMonthAndAddToMap(
@@ -224,9 +228,9 @@ public final class IcannReportingUploadAction implements Runnable {
     return cursors.build();
   }
 
-  private ImmutableMap<Key<Cursor>, Registry> loadKeyMap(
+  private ImmutableMap<VKey<? extends Cursor>, Registry> loadKeyMap(
       ImmutableSet<Registry> registries, CursorType type) {
-    return Maps.uniqueIndex(registries, r -> Cursor.createKey(type, r));
+    return Maps.uniqueIndex(registries, r -> Cursor.createVKey(type, r.getTldStr()));
   }
 
   /**
@@ -235,7 +239,9 @@ public final class IcannReportingUploadAction implements Runnable {
    * next month.
    */
   private ImmutableMap<Cursor, String> defaultNullCursorsToNextMonthAndAddToMap(
-      Map<Key<Cursor>, Registry> keyMap, CursorType type, Map<Key<Cursor>, Cursor> cursorMap) {
+      Map<VKey<? extends Cursor>, Registry> keyMap,
+      CursorType type,
+      Map<VKey<? extends Cursor>, Cursor> cursorMap) {
     ImmutableMap.Builder<Cursor, String> cursors = new ImmutableMap.Builder<>();
     keyMap.forEach(
         (key, registry) -> {
@@ -249,7 +255,7 @@ public final class IcannReportingUploadAction implements Runnable {
                       clock.nowUtc().withDayOfMonth(1).withTimeAtStartOfDay().plusMonths(1),
                       registry));
           if (!cursorMap.containsValue(cursor)) {
-            tm().transact(() -> ofy().save().entity(cursor));
+            tm().put(cursor);
           }
           cursors.put(cursor, registry.getTldStr());
         });
@@ -279,6 +285,7 @@ public final class IcannReportingUploadAction implements Runnable {
         String.format(
             "Report Filename - Upload status:\n%s",
             reportSummary.entrySet().stream()
+                .sorted(Map.Entry.comparingByKey())
                 .map(
                     (e) ->
                         String.format("%s - %s", e.getKey(), e.getValue() ? "SUCCESS" : "FAILURE"))

core/src/test/java/google/registry/export/SyncGroupMembersActionTest.java

@@ -16,10 +16,10 @@ package google.registry.export;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.export.SyncGroupMembersAction.getGroupEmailAddressForContactType;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registrar.RegistrarContact.Type.ADMIN;
 import static google.registry.model.registrar.RegistrarContact.Type.MARKETING;
 import static google.registry.model.registrar.RegistrarContact.Type.TECH;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
@@ -40,12 +40,13 @@ import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.request.Response;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
 import google.registry.util.Retrier;
 import java.io.IOException;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /**
@@ -54,6 +55,7 @@ import org.junit.jupiter.api.extension.RegisterExtension;
  * <p>Note that this relies on the registrars NewRegistrar and TheRegistrar created by default in
  * {@link AppEngineExtension}.
  */
+@DualDatabaseTest
 public class SyncGroupMembersActionTest {
 
   @RegisterExtension
@@ -74,7 +76,7 @@ public class SyncGroupMembersActionTest {
     action.run();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getGroupEmailAddressForContactType_convertsToLowercase() {
     assertThat(getGroupEmailAddressForContactType(
             "SomeRegistrar",
@@ -83,7 +85,7 @@ public class SyncGroupMembersActionTest {
         .isEqualTo("someregistrar-primary-contacts@domain-registry.example");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getGroupEmailAddressForContactType_convertsNonAlphanumericChars() {
     assertThat(getGroupEmailAddressForContactType(
             "Weird.ಠ_ಠRegistrar",
@@ -92,7 +94,7 @@ public class SyncGroupMembersActionTest {
         .isEqualTo("weirdregistrar-marketing-contacts@domain-registry.example");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_doPost_noneModified() {
     persistResource(
         loadRegistrar("NewRegistrar").asBuilder().setContactsRequireSyncing(false).build());
@@ -105,7 +107,7 @@ public class SyncGroupMembersActionTest {
     assertThat(loadRegistrar("NewRegistrar").getContactsRequireSyncing()).isFalse();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_doPost_syncsNewContact() throws Exception {
     runAction();
     verify(connection).addMemberToGroup(
@@ -117,7 +119,7 @@ public class SyncGroupMembersActionTest {
     assertThat(loadRegistrar("NewRegistrar").getContactsRequireSyncing()).isFalse();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_doPost_removesOldContact() throws Exception {
     when(connection.getMembersOfGroup("newregistrar-primary-contacts@domain-registry.example"))
         .thenReturn(ImmutableSet.of("defunct@example.com", "janedoe@theregistrar.com"));
@@ -128,13 +130,12 @@ public class SyncGroupMembersActionTest {
     assertThat(loadRegistrar("NewRegistrar").getContactsRequireSyncing()).isFalse();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_doPost_removesAllContactsFromGroup() throws Exception {
     when(connection.getMembersOfGroup("newregistrar-primary-contacts@domain-registry.example"))
         .thenReturn(ImmutableSet.of("defunct@example.com", "janedoe@theregistrar.com"));
-    ofy().deleteWithoutBackup()
+    tm().transact(
-        .entities(loadRegistrar("NewRegistrar").getContacts())
+            () -> loadRegistrar("NewRegistrar").getContacts().forEach(tm()::deleteWithoutBackup));
-        .now();
     runAction();
     verify(connection).removeMemberFromGroup(
         "newregistrar-primary-contacts@domain-registry.example", "defunct@example.com");
@@ -144,7 +145,7 @@ public class SyncGroupMembersActionTest {
     assertThat(loadRegistrar("NewRegistrar").getContactsRequireSyncing()).isFalse();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_doPost_addsAndRemovesContacts_acrossMultipleRegistrars() throws Exception {
     when(connection.getMembersOfGroup("newregistrar-primary-contacts@domain-registry.example"))
         .thenReturn(ImmutableSet.of("defunct@example.com", "janedoe@theregistrar.com"));
@@ -192,7 +193,7 @@ public class SyncGroupMembersActionTest {
         .isEmpty();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_doPost_gracefullyHandlesExceptionForSingleRegistrar() throws Exception {
     when(connection.getMembersOfGroup("newregistrar-primary-contacts@domain-registry.example"))
         .thenReturn(ImmutableSet.of());
@@ -211,7 +212,7 @@ public class SyncGroupMembersActionTest {
     assertThat(loadRegistrar("TheRegistrar").getContactsRequireSyncing()).isTrue();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_doPost_retriesOnTransientException() throws Exception {
     doThrow(IOException.class)
         .doNothing()

core/src/test/java/google/registry/export/sheet/SyncRegistrarsSheetTest.java

@@ -18,8 +18,9 @@ import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.config.RegistryConfig.getDefaultRegistrarWhoisServer;
 import static google.registry.model.common.Cursor.CursorType.SYNC_REGISTRAR_SHEET;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadByKey;
 import static google.registry.testing.DatabaseHelper.persistNewRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.DatabaseHelper.persistSimpleResources;
@@ -40,11 +41,12 @@ import google.registry.model.registrar.RegistrarAddress;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatabaseHelper;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
+import google.registry.testing.TestOfyAndSql;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.mockito.ArgumentCaptor;
@@ -54,6 +56,7 @@ import org.mockito.junit.jupiter.MockitoExtension;
 
 /** Unit tests for {@link SyncRegistrarsSheet}. */
 @ExtendWith(MockitoExtension.class)
+@DualDatabaseTest
 public class SyncRegistrarsSheetTest {
 
   @RegisterExtension
@@ -78,16 +81,22 @@ public class SyncRegistrarsSheetTest {
   void beforeEach() {
     inject.setStaticField(Ofy.class, "clock", clock);
     createTld("example");
-    // Remove Registrar entities created by AppEngineRule.
+    // Remove Registrar entities created by AppEngineRule (and RegistrarContact's, for jpa).
+    // We don't do this for ofy because ofy's loadAllOf() can't be called in a transaction but
+    // _must_ be called in a transaction in JPA.
+    if (!tm().isOfy()) {
+      tm().transact(() -> tm().loadAllOf(RegistrarContact.class))
+          .forEach(DatabaseHelper::deleteResource);
+    }
     Registrar.loadAll().forEach(DatabaseHelper::deleteResource);
   }
 
-  @Test
+  @TestOfyAndSql
   void test_wereRegistrarsModified_noRegistrars_returnsFalse() {
     assertThat(newSyncRegistrarsSheet().wereRegistrarsModified()).isFalse();
   }
 
-  @Test
+  @TestOfyAndSql
   void test_wereRegistrarsModified_atDifferentCursorTimes() {
     persistNewRegistrar("SomeRegistrar", "Some Registrar Inc.", Registrar.Type.REAL, 8L);
     persistResource(Cursor.createGlobal(SYNC_REGISTRAR_SHEET, clock.nowUtc().minusHours(1)));
@@ -96,9 +105,8 @@ public class SyncRegistrarsSheetTest {
     assertThat(newSyncRegistrarsSheet().wereRegistrarsModified()).isFalse();
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun() throws Exception {
-    DateTime beforeExecution = clock.nowUtc();
     persistResource(new Registrar.Builder()
         .setClientId("anotherregistrar")
         .setRegistrarName("Another Registrar LLC")
@@ -180,8 +188,8 @@ public class SyncRegistrarsSheetTest {
             .setTypes(ImmutableSet.of(RegistrarContact.Type.TECH))
         .build());
     // Use registrar key for contacts' parent.
+    DateTime registrarCreationTime = persistResource(registrar).getCreationTime();
     persistSimpleResources(contacts);
-    persistResource(registrar);
 
     clock.advanceBy(standardMinutes(1));
     newSyncRegistrarsSheet().run("foobar");
@@ -275,8 +283,8 @@ public class SyncRegistrarsSheetTest {
     assertThat(row).containsEntry("address.countryCode", "US");
     assertThat(row).containsEntry("phoneNumber", "+1.2223334444");
     assertThat(row).containsEntry("faxNumber", "");
-    assertThat(row.get("creationTime")).isEqualTo(beforeExecution.toString());
+    assertThat(row.get("creationTime")).isEqualTo(registrarCreationTime.toString());
-    assertThat(row.get("lastUpdateTime")).isEqualTo(beforeExecution.toString());
+    assertThat(row.get("lastUpdateTime")).isEqualTo(registrarCreationTime.toString());
     assertThat(row).containsEntry("allowedTlds", "example");
     assertThat(row).containsEntry("blockPremiumNames", "false");
     assertThat(row).containsEntry("ipAddressAllowList", "");
@@ -309,8 +317,8 @@ public class SyncRegistrarsSheetTest {
     assertThat(row).containsEntry("address.countryCode", "US");
     assertThat(row).containsEntry("phoneNumber", "+1.2125551212");
     assertThat(row).containsEntry("faxNumber", "+1.2125551213");
-    assertThat(row.get("creationTime")).isEqualTo(beforeExecution.toString());
+    assertThat(row.get("creationTime")).isEqualTo(registrarCreationTime.toString());
-    assertThat(row.get("lastUpdateTime")).isEqualTo(beforeExecution.toString());
+    assertThat(row.get("lastUpdateTime")).isEqualTo(registrarCreationTime.toString());
     assertThat(row).containsEntry("allowedTlds", "");
     assertThat(row).containsEntry("whoisServer", "whois.example.com");
     assertThat(row).containsEntry("blockPremiumNames", "false");
@@ -320,12 +328,12 @@ public class SyncRegistrarsSheetTest {
     assertThat(row).containsEntry("icannReferralEmail", "jim@example.net");
     assertThat(row).containsEntry("billingAccountMap", "{}");
 
-    Cursor cursor = ofy().load().key(Cursor.createGlobalKey(SYNC_REGISTRAR_SHEET)).now();
+    Cursor cursor = loadByKey(Cursor.createGlobalVKey(SYNC_REGISTRAR_SHEET));
     assertThat(cursor).isNotNull();
-    assertThat(cursor.getCursorTime()).isGreaterThan(beforeExecution);
+    assertThat(cursor.getCursorTime()).isGreaterThan(registrarCreationTime);
   }
 
-  @Test
+  @TestOfyAndSql
   void testRun_missingValues_stillWorks() throws Exception {
     persistNewRegistrar("SomeRegistrar", "Some Registrar", Registrar.Type.REAL, 8L);
 

core/src/test/java/google/registry/reporting/icann/IcannReportingUploadActionTest.java

@@ -15,8 +15,9 @@
 package google.registry.reporting.icann;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTlds;
+import static google.registry.testing.DatabaseHelper.loadByKey;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.GcsTestingUtils.writeGcsFile;
 import static google.registry.testing.LogsSubject.assertAboutLogs;
@@ -38,10 +39,12 @@ import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.registry.Registry;
 import google.registry.request.HttpException.ServiceUnavailableException;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeLockHandler;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.FakeSleeper;
+import google.registry.testing.TestOfyAndSql;
 import google.registry.util.EmailMessage;
 import google.registry.util.Retrier;
 import google.registry.util.SendEmailService;
@@ -51,10 +54,10 @@ import java.util.logging.Logger;
 import javax.mail.internet.InternetAddress;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link google.registry.reporting.icann.IcannReportingUploadAction} */
+@DualDatabaseTest
 class IcannReportingUploadActionTest {
 
   @RegisterExtension
@@ -126,7 +129,7 @@ class IcannReportingUploadActionTest {
     loggerToIntercept.addHandler(logHandler);
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess() throws Exception {
     IcannReportingUploadAction action = createAction();
     action.run();
@@ -142,14 +145,14 @@ class IcannReportingUploadActionTest {
                 "ICANN Monthly report upload summary: 3/4 succeeded",
                 "Report Filename - Upload status:\n"
                     + "foo-activity-200606.csv - SUCCESS\n"
-                    + "tld-activity-200606.csv - FAILURE\n"
                     + "foo-transactions-200606.csv - SUCCESS\n"
+                    + "tld-activity-200606.csv - FAILURE\n"
                     + "tld-transactions-200606.csv - SUCCESS",
                 new InternetAddress("recipient@example.com"),
                 new InternetAddress("sender@example.com")));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_january() throws Exception {
     clock.setTo(DateTime.parse("2006-01-22T00:30:00Z"));
     persistResource(
@@ -186,7 +189,7 @@ class IcannReportingUploadActionTest {
                 new InternetAddress("sender@example.com")));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_advancesCursor() throws Exception {
     writeGcsFile(
         gcsService,
@@ -195,21 +198,17 @@ class IcannReportingUploadActionTest {
     when(mockReporter.send(PAYLOAD_SUCCESS, "tld-activity-200606.csv")).thenReturn(true);
     IcannReportingUploadAction action = createAction();
     action.run();
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
-    Cursor cursor =
+    Cursor cursor = loadByKey(Cursor.createVKey(CursorType.ICANN_UPLOAD_ACTIVITY, "tld"));
-        ofy()
-            .load()
-            .key(Cursor.createKey(CursorType.ICANN_UPLOAD_ACTIVITY, Registry.get("tld")))
-            .now();
     assertThat(cursor.getCursorTime()).isEqualTo(DateTime.parse("2006-08-01TZ"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_noUploadsNeeded() throws Exception {
     clock.setTo(DateTime.parse("2006-5-01T00:30:00Z"));
     IcannReportingUploadAction action = createAction();
     action.run();
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
     verifyNoMoreInteractions(mockReporter);
     verify(emailService)
         .sendEmail(
@@ -220,7 +219,7 @@ class IcannReportingUploadActionTest {
                 new InternetAddress("sender@example.com")));
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_withRetry() throws Exception {
     IcannReportingUploadAction action = createAction();
     when(mockReporter.send(PAYLOAD_SUCCESS, "tld-transactions-200606.csv"))
@@ -238,14 +237,14 @@ class IcannReportingUploadActionTest {
                 "ICANN Monthly report upload summary: 3/4 succeeded",
                 "Report Filename - Upload status:\n"
                     + "foo-activity-200606.csv - SUCCESS\n"
-                    + "tld-activity-200606.csv - FAILURE\n"
                     + "foo-transactions-200606.csv - SUCCESS\n"
+                    + "tld-activity-200606.csv - FAILURE\n"
                     + "tld-transactions-200606.csv - SUCCESS",
                 new InternetAddress("recipient@example.com"),
                 new InternetAddress("sender@example.com")));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_quicklySkipsOverNonRetryableUploadException() throws Exception {
     runTest_nonRetryableException(
         new IOException(
@@ -253,36 +252,28 @@ class IcannReportingUploadActionTest {
                 + " passed.</msg>"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_quicklySkipsOverIpAllowListException() throws Exception {
     runTest_nonRetryableException(
         new IOException("Your IP address 25.147.130.158 is not allowed to connect"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_cursorIsNotAdvancedForward() throws Exception {
     runTest_nonRetryableException(
         new IOException("Your IP address 25.147.130.158 is not allowed to connect"));
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
-    Cursor cursor =
+    Cursor cursor = loadByKey(Cursor.createVKey(CursorType.ICANN_UPLOAD_ACTIVITY, "tld"));
-        ofy()
-            .load()
-            .key(Cursor.createKey(CursorType.ICANN_UPLOAD_ACTIVITY, Registry.get("tld")))
-            .now();
     assertThat(cursor.getCursorTime()).isEqualTo(DateTime.parse("2006-07-01TZ"));
   }
 
-  @Test
+  @TestOfyAndSql
   void testNotRunIfCursorDateIsAfterToday() throws Exception {
     clock.setTo(DateTime.parse("2006-05-01T00:30:00Z"));
     IcannReportingUploadAction action = createAction();
     action.run();
-    ofy().clearSessionCache();
+    tm().clearSessionCache();
-    Cursor cursor =
+    Cursor cursor = loadByKey(Cursor.createVKey(CursorType.ICANN_UPLOAD_ACTIVITY, "foo"));
-        ofy()
-            .load()
-            .key(Cursor.createKey(CursorType.ICANN_UPLOAD_ACTIVITY, Registry.get("foo")))
-            .now();
     assertThat(cursor.getCursorTime()).isEqualTo(DateTime.parse("2006-07-01TZ"));
     verifyNoMoreInteractions(mockReporter);
   }
@@ -306,14 +297,14 @@ class IcannReportingUploadActionTest {
                 "ICANN Monthly report upload summary: 3/4 succeeded",
                 "Report Filename - Upload status:\n"
                     + "foo-activity-200606.csv - SUCCESS\n"
-                    + "tld-activity-200606.csv - FAILURE\n"
                     + "foo-transactions-200606.csv - SUCCESS\n"
+                    + "tld-activity-200606.csv - FAILURE\n"
                     + "tld-transactions-200606.csv - SUCCESS",
                 new InternetAddress("recipient@example.com"),
                 new InternetAddress("sender@example.com")));
   }
 
-  @Test
+  @TestOfyAndSql
   void testFail_fileNotFound() throws Exception {
     clock.setTo(DateTime.parse("2006-01-22T00:30:00Z"));
     persistResource(
@@ -329,7 +320,7 @@ class IcannReportingUploadActionTest {
                 + " tld-activity-200512.csv did not exist");
   }
 
-  @Test
+  @TestOfyAndSql
   void testWarning_fileNotStagedYet() throws Exception {
     persistResource(
         Cursor.create(
@@ -346,7 +337,7 @@ class IcannReportingUploadActionTest {
                 + " yet.");
   }
 
-  @Test
+  @TestOfyAndSql
   void testFailure_lockIsntAvailable() throws Exception {
     IcannReportingUploadAction action = createAction();
     action.lockHandler = new FakeLockHandler(false);
@@ -357,7 +348,7 @@ class IcannReportingUploadActionTest {
         .contains("Lock for IcannReportingUploadAction already in use");
   }
 
-  @Test
+  @TestOfyAndSql
   void testSuccess_nullCursorsInitiatedToFirstOfNextMonth() throws Exception {
     createTlds("new");
 
@@ -375,21 +366,16 @@ class IcannReportingUploadActionTest {
                 "ICANN Monthly report upload summary: 3/4 succeeded",
                 "Report Filename - Upload status:\n"
                     + "foo-activity-200606.csv - SUCCESS\n"
-                    + "tld-activity-200606.csv - FAILURE\n"
                     + "foo-transactions-200606.csv - SUCCESS\n"
+                    + "tld-activity-200606.csv - FAILURE\n"
                     + "tld-transactions-200606.csv - SUCCESS",
                 new InternetAddress("recipient@example.com"),
                 new InternetAddress("sender@example.com")));
 
     Cursor newActivityCursor =
-        ofy()
+        loadByKey(Cursor.createVKey(CursorType.ICANN_UPLOAD_ACTIVITY, "new"));
-            .load()
-            .key(Cursor.createKey(CursorType.ICANN_UPLOAD_ACTIVITY, Registry.get("new")))
-            .now();
     assertThat(newActivityCursor.getCursorTime()).isEqualTo(DateTime.parse("2006-08-01TZ"));
-    Cursor newTransactionCursor =
+    Cursor newTransactionCursor = loadByKey(Cursor.createVKey(CursorType.ICANN_UPLOAD_TX, "new"));
-        ofy().load().key(Cursor.createKey(CursorType.ICANN_UPLOAD_TX, Registry.get("new"))).now();
     assertThat(newTransactionCursor.getCursorTime()).isEqualTo(DateTime.parse("2006-08-01TZ"));
   }
 }
-