Remove some unnecessary Ofy key creation (#1212)

core/src/test/java/google/registry/beam/common/RegistryJpaReadTest.java

@@ -45,7 +45,6 @@ import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatabaseHelper;
 import google.registry.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
-import google.registry.testing.SystemPropertyExtension;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.values.PCollection;
 import org.joda.time.DateTime;
@@ -65,14 +64,6 @@ public class RegistryJpaReadTest {
   @Order(Order.DEFAULT - 1)
   final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
 
-  // The pipeline runner on Kokoro sometimes mistakes the platform as appengine, resulting in
-  // a null thread factory. The cause is unknown but it may be due to the interaction with
-  // the DatastoreEntityExtension above. To work around the problem, we explicitly unset the
-  // relevant property before test starts.
-  @RegisterExtension
-  final transient SystemPropertyExtension systemPropertyExtension =
-      new SystemPropertyExtension().setProperty("com.google.appengine.runtime.environment", null);
-
   @RegisterExtension
   final transient JpaIntegrationTestExtension database =
       new JpaTestRules.Builder().withClock(fakeClock).buildIntegrationTestRule();

core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java

@@ -18,6 +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.tm;
 import static google.registry.testing.AppEngineExtension.makeRegistrar1;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -49,7 +50,6 @@ import google.registry.model.reporting.Spec11ThreatMatchDao;
 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.testing.DatastoreEntityExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
@@ -71,7 +71,6 @@ import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
 import org.json.JSONObject;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -113,8 +112,6 @@ class Spec11PipelineTest {
           ThreatMatch.create("THREAT_TYPE_UNSPECIFIED", "no-eamil.com"),
           ThreatMatch.create("UNWANTED_SOFTWARE", "anti-anti-anti-virus.dev"));
 
-  // This extension is only needed because Spec11ThreatMatch uses Ofy to generate the ID. Can be
-  // removed after the SQL migration.
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
   final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
@@ -136,12 +133,9 @@ class Spec11PipelineTest {
   private PCollection<KV<Subdomain, ThreatMatch>> threatMatches;
 
   ImmutableSet<Spec11ThreatMatch> sqlThreatMatches;
-  TransactionManager tm;
 
   @BeforeEach
   void beforeEach() throws Exception {
-    tm = tm();
-    TransactionManagerFactory.setTm(jpaTm());
     reportingBucketUrl = Files.createDirectory(tmpDir.resolve(REPORTING_BUCKET_URL)).toFile();
     options.setDate(DATE);
     options.setSafeBrowsingApiKey(SAFE_BROWSING_API_KEY);
@@ -196,11 +190,6 @@ class Spec11PipelineTest {
                 .build());
   }
 
-  @AfterEach
-  void afterEach() {
-    TransactionManagerFactory.setTm(tm);
-  }
-
   @Test
   void testSuccess_fullSqlPipeline() throws Exception {
     setupCloudSql();
@@ -241,6 +230,8 @@ class Spec11PipelineTest {
   }
 
   private void setupCloudSql() {
+    TransactionManager originalTm = tm();
+    setTm(jpaTm());
     persistNewRegistrar("TheRegistrar");
     persistNewRegistrar("NewRegistrar");
     Registrar registrar1 =
@@ -280,6 +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);
   }
 
   private void verifySaveToGcs() throws Exception {

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

@@ -44,6 +44,7 @@ import com.google.common.collect.Sets;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.ObjectifyFilter;
+import google.registry.model.IdService;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
@@ -466,7 +467,7 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
     if (withDatastore) {
       ObjectifyService.initOfy();
       // Reset id allocation in ObjectifyService so that ids are deterministic in tests.
-      ObjectifyService.resetNextTestId();
+      IdService.resetNextTestId();
       this.ofyTestEntities.forEach(AppEngineExtension::register);
     }
   }

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

@@ -28,10 +28,10 @@ import static google.registry.config.RegistryConfig.getContactAndHostRoidSuffix;
 import static google.registry.config.RegistryConfig.getContactAutomaticTransferLength;
 import static google.registry.model.EppResourceUtils.createDomainRepoId;
 import static google.registry.model.EppResourceUtils.createRepoId;
+import static google.registry.model.IdService.allocateId;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
 import static google.registry.model.ResourceTransferUtils.createTransferResponse;
-import static google.registry.model.ofy.ObjectifyService.allocateId;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
@@ -90,7 +90,6 @@ import google.registry.model.host.HostResource;
 import google.registry.model.index.EppResourceIndex;
 import google.registry.model.index.EppResourceIndexBucket;
 import google.registry.model.index.ForeignKeyIndex;
-import google.registry.model.ofy.ObjectifyService;
 import google.registry.model.poll.PollMessage;
 import google.registry.model.pricing.StaticPremiumListPricingEngine;
 import google.registry.model.registrar.Registrar;
@@ -957,12 +956,12 @@ public class DatabaseHelper {
 
   /** Returns a newly allocated, globally unique domain repoId of the format HEX-TLD. */
   public static String generateNewDomainRoid(String tld) {
-    return createDomainRepoId(ObjectifyService.allocateId(), tld);
+    return createDomainRepoId(allocateId(), tld);
   }
 
   /** Returns a newly allocated, globally unique contact/host repoId of the format HEX_TLD-ROID. */
   public static String generateNewContactHostRoid() {
-    return createRepoId(ObjectifyService.allocateId(), getContactAndHostRoidSuffix());
+    return createRepoId(allocateId(), getContactAndHostRoidSuffix());
   }
 
   /**

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

@@ -14,6 +14,8 @@
 
 package google.registry.testing;
 
+import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+
 import com.google.apphosting.api.ApiProxy;
 import com.google.apphosting.api.ApiProxy.Environment;
 import java.util.Map;
@@ -24,7 +26,11 @@ import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
 
 /**
  * Allows instantiation of Datastore {@code Entity entities} without the heavyweight {@code
- * AppEngineRule}.
+ * AppEngineExtension}.
+ *
+ * <p>When an Ofy key is created, by calling the various Key.create() methods, whether the current
+ * executing thread is a GAE thread is checked, which this extension masquerades. This happens
+ * frequently when an entity which has the key of another entity as a field is instantiated.
  *
  * <p>When used together with {@code JpaIntegrationWithCoverageExtension} or @{@code
  * TestPipelineExtension}, this extension must be registered first. For consistency's sake, it is
@@ -41,6 +47,10 @@ public class DatastoreEntityExtension implements BeforeEachCallback, AfterEachCa
   @Override
   public void beforeEach(ExtensionContext context) {
     ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
+    // In order to create keys for entities they must be registered with Ofy. Calling this method
+    // will load the ObjectifyService class, whose static initialization block registers all Ofy
+    // entities.
+    auditedOfy();
   }
 
   @Override