Fix flaky tests due to Entity name conflicts (#569)

* Fix flaky tests due to Entity name conflicts

Objectify siliently replaces current registration of a given kind
when another class is registered for this kind. There are
several TestObject classes in the current code base, which by
default are all mapped to the same kind.

Tests have only been flaky because impacted tests need to run
in specific orders for failures to happen. Using multiple executors
in Gradle also reduced the likely hood of errors. To reproduce the
problem run the following tests in order (e.g., by putting them in
a test suite):
1. ExportCommitLogDiffActionTest
2. CreateAutoTimestampTest
3. RestoreCommitLogsActionTest

In this PR, we
- Made sure all entities have unique kinds.
- Made all test entities register with AppEngineRule instead of directly
  with ObjectifyService.
- Added code in AppEngineRule to check for re-registrations.
- Added presumit check for forbidden direct registration.

config/presubmits.py

@@ -99,6 +99,15 @@ PRESUBMITS = {
         "System.(out|err).println is only allowed in tools/ packages. Please "
         "use a logger instead.",
 
+    # ObjectifyService.register is restricted to main/ or AppEngineRule.
+    PresubmitCheck(
+        r".*\bObjectifyService\.register", "java", {
+            "/build/", "/generated/", "node_modules/", "src/main/",
+            "AppEngineRule.java"
+        }):
+      "ObjectifyService.register is not allowed in tests. Please use "
+      "AppengineRule.register instead.",
+
     # PostgreSQLContainer instantiation must specify docker tag
     PresubmitCheck(
         r"[\s\S]*new\s+PostgreSQLContainer(<[\s\S]*>)?\(\s*\)[\s\S]*",

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

@@ -151,11 +151,14 @@ public class ObjectifyService {
       String kind = Key.getKind(clazz);
       boolean registered = factory().getMetadata(kind) != null;
       if (clazz.isAnnotationPresent(Entity.class)) {
-        // Objectify silently ignores re-registrations for a given kind string, even if the classes
-        // being registered are distinct. Throw an exception if that would happen here.
-        checkState(!registered,
+        // Objectify silently replaces current registration for a given kind string when a different
+        // class is registered again for this kind. For simplicity's sake, throw an exception on any
+        // re-registration.
+        checkState(
+            !registered,
             "Kind '%s' already registered, cannot register new @Entity %s",
-            kind, clazz.getCanonicalName());
+            kind,
+            clazz.getCanonicalName());
       } else if (clazz.isAnnotationPresent(EntitySubclass.class)) {
         // Ensure that any @EntitySubclass classes have also had their parent @Entity registered,
         // which Objectify nominally requires but doesn't enforce in 4.x (though it may in 5.x).

core/src/test/java/google/registry/backup/ExportCommitLogDiffActionTest.java

@@ -50,7 +50,11 @@ import org.junit.runners.JUnit4;
 public class ExportCommitLogDiffActionTest {
 
   @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
+  public final AppEngineRule appEngine =
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(TestObject.class)
+          .build();
 
   /** Local GCS service available for testing. */
   private final GcsService gcsService = GcsServiceFactory.createGcsService();

core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java

@@ -71,7 +71,11 @@ public class RestoreCommitLogsActionTest {
   final GcsService gcsService = createGcsService();
 
   @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
+  public final AppEngineRule appEngine =
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(TestObject.class)
+          .build();
 
   @Before
   public void init() {

core/src/test/java/google/registry/model/CreateAutoTimestampTest.java

@@ -19,12 +19,10 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static org.joda.time.DateTimeZone.UTC;
 
-import com.googlecode.objectify.ObjectifyService;
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.common.CrossTldSingleton;
 import google.registry.testing.AppEngineRule;
 import org.joda.time.DateTime;
-import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -35,19 +33,18 @@ import org.junit.runners.JUnit4;
 public class CreateAutoTimestampTest {
 
   @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
+  public final AppEngineRule appEngine =
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(TestObject.class)
+          .build();
 
   /** Timestamped class. */
-  @Entity
+  @Entity(name = "CatTestEntity")
   public static class TestObject extends CrossTldSingleton {
     CreateAutoTimestamp createTime = CreateAutoTimestamp.create(null);
   }
 
-  @Before
-  public void before() {
-    ObjectifyService.register(TestObject.class);
-  }
-
   private TestObject reload() {
     return ofy().load().entity(new TestObject()).now();
   }

core/src/test/java/google/registry/model/ImmutableObjectTest.java

@@ -27,7 +27,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.ObjectifyService;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import google.registry.testing.AppEngineRule;
@@ -39,7 +38,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import org.joda.time.DateTime;
-import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -50,12 +48,11 @@ import org.junit.runners.JUnit4;
 public class ImmutableObjectTest {
 
   @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
-
-  @Before
-  public void register() {
-    ObjectifyService.register(ValueObject.class);
-  }
+  public final AppEngineRule appEngine =
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(ValueObject.class)
+          .build();
 
   /** Simple subclass of ImmutableObject. */
   public static class SimpleObject extends ImmutableObject {

core/src/test/java/google/registry/model/UpdateAutoTimestampTest.java

@@ -19,12 +19,10 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static org.joda.time.DateTimeZone.UTC;
 
-import com.googlecode.objectify.ObjectifyService;
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.common.CrossTldSingleton;
 import google.registry.testing.AppEngineRule;
 import org.joda.time.DateTime;
-import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -35,19 +33,18 @@ import org.junit.runners.JUnit4;
 public class UpdateAutoTimestampTest {
 
   @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
+  public final AppEngineRule appEngine =
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(TestObject.class)
+          .build();
 
   /** Timestamped class. */
-  @Entity
+  @Entity(name = "UatTestEntity")
   public static class TestObject extends CrossTldSingleton {
     UpdateAutoTimestamp updateTime = UpdateAutoTimestamp.create(null);
   }
 
-  @Before
-  public void before() {
-    ObjectifyService.register(TestObject.class);
-  }
-
   private TestObject reload() {
     return ofy().load().entity(new TestObject()).now();
   }

core/src/test/java/google/registry/model/ofy/OfyCommitLogTest.java

@@ -16,7 +16,6 @@ package google.registry.model.ofy;
 
 import static com.google.appengine.api.datastore.EntityTranslator.convertToPb;
 import static com.google.common.truth.Truth.assertThat;
-import static com.googlecode.objectify.ObjectifyService.register;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.CommitLogBucket.getBucketKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
@@ -47,7 +46,11 @@ import org.junit.runners.JUnit4;
 public class OfyCommitLogTest {
 
   @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
+  public final AppEngineRule appEngine =
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(TestVirtualObject.class, Root.class, Child.class)
+          .build();
 
   @Rule
   public final InjectRule inject = new InjectRule();
@@ -56,8 +59,6 @@ public class OfyCommitLogTest {
 
   @Before
   public void before() {
-    register(Root.class);
-    register(Child.class);
     inject.setStaticField(Ofy.class, "clock", clock);
   }
 

core/src/test/java/google/registry/model/translators/CommitLogRevisionsTranslatorFactoryTest.java

@@ -23,7 +23,6 @@ import static org.joda.time.Duration.standardHours;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSortedMap;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.ObjectifyService;
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.common.CrossTldSingleton;
 import google.registry.model.ofy.CommitLogManifest;
@@ -45,13 +44,17 @@ public class CommitLogRevisionsTranslatorFactoryTest {
 
   private static final DateTime START_TIME = DateTime.parse("2000-01-01TZ");
 
-  @Entity
+  @Entity(name = "ClrtfTestEntity")
   public static class TestObject extends CrossTldSingleton {
     ImmutableSortedMap<DateTime, Key<CommitLogManifest>> revisions = ImmutableSortedMap.of();
   }
 
   @Rule
-  public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
+  public final AppEngineRule appEngine =
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(TestObject.class)
+          .build();
 
   @Rule
   public final InjectRule inject = new InjectRule();
@@ -60,7 +63,6 @@ public class CommitLogRevisionsTranslatorFactoryTest {
 
   @Before
   public void before() {
-    ObjectifyService.register(TestObject.class);
     inject.setStaticField(Ofy.class, "clock", clock);
   }
 

core/src/test/java/google/registry/persistence/VKeyTest.java

@@ -28,7 +28,10 @@ public class VKeyTest {
 
   @Rule
   public final AppEngineRule appEngineRule =
-      AppEngineRule.builder().withDatastoreAndCloudSql().build();
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .withOfyTestEntities(TestObject.class)
+          .build();
 
   public VKeyTest() {}
 

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

@@ -34,6 +34,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 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.ofy.ObjectifyService;
 import google.registry.model.registrar.Registrar;
@@ -50,6 +51,7 @@ import java.io.File;
 import java.io.IOException;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.logging.LogManager;
 import javax.annotation.Nullable;
@@ -127,10 +129,14 @@ public final class AppEngineRule extends ExternalResource
   private String taskQueueXml;
   private UserInfo userInfo;
 
+  // Test Objectify entity classes to be used with this AppEngineRule instance.
+  private ImmutableList<Class<?>> ofyTestEntities;
+
   /** Builder for {@link AppEngineRule}. */
   public static class Builder {
 
     private AppEngineRule rule = new AppEngineRule();
+    private ImmutableList.Builder<Class<?>> ofyTestEntities = new ImmutableList.Builder();
 
     /** Turn on the Datastore service and the Cloud SQL service. */
     public Builder withDatastoreAndCloudSql() {
@@ -181,10 +187,29 @@ public final class AppEngineRule extends ExternalResource
       return this;
     }
 
+    /**
+     * Declares test-only entities to be registered with {@code ObjectifyService}.
+     *
+     * <p>Note that {@code ObjectifyService} silently replaces the current registration for a given
+     * kind when a different class is registered for this kind. Since {@code ObjectifyService} does
+     * not support de-registration, each test entity class must be of a unique kind across the
+     * entire code base. Although this requirement can be worked around by using different {@code
+     * ObjectifyService} instances for each test (class), the setup overhead would rise
+     * significantly.
+     *
+     * @see AppEngineRule#register(Class)
+     */
+    @SafeVarargs
+    public final Builder withOfyTestEntities(Class<?>... entities) {
+      ofyTestEntities.add(entities);
+      return this;
+    }
+
     public AppEngineRule build() {
       checkState(
           !rule.enableJpaEntityCoverageCheck || rule.withDatastoreAndCloudSql,
           "withJpaEntityCoverageCheck enabled without Cloud SQL");
+      rule.ofyTestEntities = this.ofyTestEntities.build();
       return rule;
     }
   }
@@ -411,6 +436,7 @@ public final class AppEngineRule extends ExternalResource
       // Reset id allocation in ObjectifyService so that ids are deterministic in tests.
       ObjectifyService.resetNextTestId();
       loadInitialData();
+      this.ofyTestEntities.forEach(AppEngineRule::register);
     }
   }
 
@@ -443,6 +469,24 @@ public final class AppEngineRule extends ExternalResource
     }
   }
 
+  /**
+   * Registers test-only Objectify entities and checks for re-registrations for the same kind by
+   * different classes.
+   */
+  private static void register(Class<?> entityClass) {
+    String kind = Key.getKind(entityClass);
+    Optional.ofNullable(com.googlecode.objectify.ObjectifyService.factory().getMetadata(kind))
+        .ifPresent(
+            meta ->
+                checkState(
+                    meta.getEntityClass() == entityClass,
+                    "Cannot register %s. The Kind %s is already registered with %s.",
+                    entityClass.getName(),
+                    kind,
+                    meta.getEntityClass().getName()));
+    com.googlecode.objectify.ObjectifyService.register(entityClass);
+  }
+
   /** Install {@code testing/logging.properties} so logging is less noisy. */
   private static void setupLogging() throws IOException {
     LogManager.getLogManager()

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

@@ -18,6 +18,7 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.base.Joiner;
+import com.googlecode.objectify.annotation.Entity;
 import java.io.File;
 import java.io.IOException;
 import org.junit.Before;
@@ -88,9 +89,29 @@ public class AppEngineRuleTest {
     assertThrows(AssertionError.class, () -> appEngineRule.after());
   }
 
+  @Test
+  public void testRegisterOfyEntities_failure() {
+    AppEngineRule appEngineRule =
+        AppEngineRule.builder()
+            .withDatastoreAndCloudSql()
+            .withOfyTestEntities(google.registry.testing.TestObject.class, TestObject.class)
+            .build();
+    String expectedErrorMessage =
+        String.format(
+            "Cannot register %s. The Kind %s is already registered with %s",
+            TestObject.class.getName(),
+            "TestObject",
+            google.registry.testing.TestObject.class.getName());
+    assertThrows(expectedErrorMessage, IllegalStateException.class, appEngineRule::before);
+    appEngineRule.after();
+  }
+
   private void writeAutoIndexFile(String content) throws IOException {
     com.google.common.io.Files.asCharSink(
             new File(temporaryFolder.getRoot(), "datastore-indexes-auto.xml"), UTF_8)
         .write(content);
   }
+
+  @Entity
+  private static final class TestObject {}
 }

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

@@ -17,7 +17,6 @@ package google.registry.testing;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.ObjectifyService;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Parent;
@@ -30,9 +29,6 @@ import google.registry.model.common.EntityGroupRoot;
  */
 @Entity
 public class TestObject extends ImmutableObject {
-  static {
-    ObjectifyService.register(TestObject.class);  // Register this kind on first reference.
-  }
 
   @Parent
   Key<EntityGroupRoot> parent;
@@ -70,9 +66,6 @@ public class TestObject extends ImmutableObject {
   @Entity
   @VirtualEntity
   public static class TestVirtualObject extends ImmutableObject {
-    static {
-      ObjectifyService.register(TestVirtualObject.class);  // Register this kind on first reference.
-    }
 
     @Id
     String id;