Streamline how to fake an App Engine environment (#1348)

Both `DatastoreEntityExtension.PlaceholderEnvironment` and `AppEngineEnvironment` does the same thing, so there is no point having both of them exist. To use `AppEngineEnvionrment` as an autoclosable requires the user to be mindful of where a fake App Engine environment is required. It is better to set this either in the `DatastoreEntityExtension` for tests, or in the worker initializer in Beam. It also makes it easier to remove the fake environment when we are completely datastore free.

Also made a change to how `IdService` allocate Ids in Beam.
<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/google/nomulus/1348)
<!-- Reviewable:end -->

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

@@ -62,7 +62,8 @@ public class RegistryJpaReadTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final transient JpaIntegrationTestExtension database =

core/src/test/java/google/registry/beam/initsql/CommitLogTransformsTest.java

@@ -71,7 +71,7 @@ class CommitLogTransformsTest implements Serializable {
   @RegisterExtension
   @Order(value = 1)
   final transient DatastoreEntityExtension datastoreEntityExtension =
-      new DatastoreEntityExtension();
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final transient TestPipelineExtension testPipeline =

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineTest.java

@@ -30,7 +30,6 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.beam.TestPipelineExtension;
 import google.registry.flows.domain.DomainFlowUtils;
 import google.registry.model.billing.BillingEvent;
@@ -106,7 +105,8 @@ class InitSqlPipelineTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension("test").allThreads(true);
 
   @RegisterExtension final transient InjectExtension injectExtension = new InjectExtension();
 
@@ -331,20 +331,18 @@ class InitSqlPipelineTest {
             .as(InitSqlPipelineOptions.class);
     InitSqlPipeline initSqlPipeline = new InitSqlPipeline(options);
     initSqlPipeline.run(testPipeline).waitUntilFinish();
-    try (AppEngineEnvironment env = new AppEngineEnvironment("test")) {
-      assertHostResourceEquals(
-          jpaTm().transact(() -> jpaTm().loadByKey(hostResource.createVKey())), hostResource);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Registrar.class)))
-          .comparingElementsUsing(immutableObjectCorrespondence("lastUpdateTime"))
-          .containsExactly(registrar1, registrar2);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class)))
-          .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
-          .containsExactly(contact1, contact2);
-      assertDomainEquals(jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey())), domain);
-      assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Cursor.class)))
-          .comparingElementsUsing(immutableObjectCorrespondence())
-          .containsExactly(globalCursor, tldCursor);
-    }
+    assertHostResourceEquals(
+        jpaTm().transact(() -> jpaTm().loadByKey(hostResource.createVKey())), hostResource);
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Registrar.class)))
+        .comparingElementsUsing(immutableObjectCorrespondence("lastUpdateTime"))
+        .containsExactly(registrar1, registrar2);
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(ContactResource.class)))
+        .comparingElementsUsing(immutableObjectCorrespondence("revisions", "updateTimestamp"))
+        .containsExactly(contact1, contact2);
+    assertDomainEquals(jpaTm().transact(() -> jpaTm().loadByKey(domain.createVKey())), domain);
+    assertThat(jpaTm().transact(() -> jpaTm().loadAllOf(Cursor.class)))
+        .comparingElementsUsing(immutableObjectCorrespondence())
+        .containsExactly(globalCursor, tldCursor);
   }
 
   private static void assertHostResourceEquals(HostResource actual, HostResource expected) {

core/src/test/java/google/registry/beam/initsql/InitSqlTestUtils.java

@@ -26,7 +26,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
 import com.google.common.truth.Truth;
 import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
-import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.VersionedEntity;
 import java.io.Serializable;
 import java.util.Collection;
@@ -124,20 +123,18 @@ public final class InitSqlTestUtils {
                       public void processElement(
                           @Element KV<String, Iterable<VersionedEntity>> input,
                           OutputReceiver<String> out) {
-                        try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-                          ImmutableList<KV<Long, Object>> actual =
-                              Streams.stream(input.getValue())
-                                  .map(InitSqlTestUtils::rawEntityToOfyWithTimestamp)
-                                  .collect(ImmutableList.toImmutableList());
-                          try {
-                            Truth.assertThat(actual)
-                                .containsExactlyElementsIn(
-                                    Stream.of(expected)
-                                        .map(InitSqlTestUtils::expectedToOfyWithTimestamp)
-                                        .collect(ImmutableList.toImmutableList()));
-                          } catch (AssertionError e) {
-                            out.output(e.toString());
-                          }
+                        ImmutableList<KV<Long, Object>> actual =
+                            Streams.stream(input.getValue())
+                                .map(InitSqlTestUtils::rawEntityToOfyWithTimestamp)
+                                .collect(ImmutableList.toImmutableList());
+                        try {
+                          Truth.assertThat(actual)
+                              .containsExactlyElementsIn(
+                                  Stream.of(expected)
+                                      .map(InitSqlTestUtils::expectedToOfyWithTimestamp)
+                                      .collect(ImmutableList.toImmutableList()));
+                        } catch (AssertionError e) {
+                          out.output(e.toString());
                         }
                       }
                     }));

core/src/test/java/google/registry/beam/initsql/LoadDatastoreSnapshotTest.java

@@ -95,7 +95,7 @@ class LoadDatastoreSnapshotTest {
   @RegisterExtension
   @Order(value = 1)
   final transient DatastoreEntityExtension datastoreEntityExtension =
-      new DatastoreEntityExtension();
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final transient TestPipelineExtension testPipeline =

core/src/test/java/google/registry/beam/invoicing/InvoicingPipelineTest.java

@@ -227,7 +227,8 @@ class InvoicingPipelineTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
 
   @RegisterExtension
   final TestPipelineExtension pipeline =

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

@@ -116,7 +116,8 @@ class Spec11PipelineTest {
 
   @RegisterExtension
   @Order(Order.DEFAULT - 1)
-  final transient DatastoreEntityExtension datastore = new DatastoreEntityExtension();
+  final transient DatastoreEntityExtension datastore =
+      new DatastoreEntityExtension().allThreads(true);
 
   @TempDir Path tmpDir;
 

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

@@ -460,7 +460,7 @@ public final class AppEngineExtension implements BeforeEachCallback, AfterEachCa
     if (withDatastore) {
       ObjectifyService.initOfy();
       // Reset id allocation in ObjectifyService so that ids are deterministic in tests.
-      IdService.resetNextTestId();
+      IdService.resetSelfAllocatedId();
       this.ofyTestEntities.forEach(AppEngineExtension::register);
     }
   }

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

@@ -14,17 +14,10 @@
 
 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.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
-import java.util.Map;
+import google.registry.model.AppEngineEnvironment;
 import org.junit.jupiter.api.extension.AfterEachCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
-import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
 
 /**
  * Allows instantiation of Datastore {@code Entity}s without the heavyweight {@link
@@ -41,13 +34,23 @@ import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
  * google.registry.model.domain.DomainBaseSqlTest} for example, and to <a
  * href="https://junit.org/junit5/docs/current/user-guide/#extensions-registration-programmatic">
  * JUnit 5 User Guide</a> for details of extension ordering.
+ *
+ * @see AppEngineEnvironment
  */
 public class DatastoreEntityExtension implements BeforeEachCallback, AfterEachCallback {
 
-  private static final Environment PLACEHOLDER_ENV = new PlaceholderEnvironment();
+  private final AppEngineEnvironment environment;
 
   private boolean allThreads = false;
 
+  public DatastoreEntityExtension(String appId) {
+    environment = new AppEngineEnvironment(appId);
+  }
+
+  public DatastoreEntityExtension() {
+    environment = new AppEngineEnvironment();
+  }
+
   /**
    * Whether all threads should be masqueraded as GAE threads.
    *
@@ -69,79 +72,19 @@ 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();
     if (allThreads) {
-      ApiProxy.setEnvironmentFactory(() -> PLACEHOLDER_ENV);
+      environment.setEnvironmentForAllThreads();
+    } else {
+      environment.setEnvironmentForCurrentThread();
     }
   }
 
   @Override
-  public void afterEach(ExtensionContext context)
-      throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
-    // Clear the cached instance.
-    ApiProxy.clearEnvironmentForCurrentThread();
+  public void afterEach(ExtensionContext context) {
     if (allThreads) {
-      Method method = ApiProxy.class.getDeclaredMethod("clearEnvironmentFactory");
-      method.setAccessible(true);
-      method.invoke(null);
-    }
-  }
-
-  private static final class PlaceholderEnvironment implements Environment {
-
-    @Override
-    public String getAppId() {
-      return "PlaceholderAppId";
-    }
-
-    @Override
-    public Map<String, Object> getAttributes() {
-      return ImmutableMap.of();
-    }
-
-    @Override
-    public String getModuleId() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public String getVersionId() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public String getEmail() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public boolean isLoggedIn() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public boolean isAdmin() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public String getAuthDomain() {
-      throw new UnsupportedOperationException();
-    }
-
-    @SuppressWarnings("deprecation")
-    @Override
-    public String getRequestNamespace() {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public long getRemainingMillis() {
-      throw new UnsupportedOperationException();
+      environment.unsetEnvironmentForAllThreads();
+    } else {
+      environment.unsetEnvironmentForCurrentThread();
     }
   }
 }