Implement DatastoreEntity/SqlEntity for many more classes (#788)

* Implement DatastoreEntity/SqlEntity for many more classes

We still have many more classes to go, but this gets us closer to
guaranteeing that we can convert from Datastore to SQL objects and back
again.

* Shift SqlEntity impl to HistoryEntry

core/src/main/java/google/registry/model/billing/BillingEvent.java

@@ -46,6 +46,7 @@ import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.transfer.TransferData.TransferServerApproveEntity;
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithLongVKey;
+import google.registry.schema.replay.DatastoreAndSqlEntity;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -274,7 +275,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "allocation_token_id")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_event_id"))
-  public static class OneTime extends BillingEvent {
+  public static class OneTime extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billable value. */
     @AttributeOverrides({
@@ -450,7 +451,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "recurrence_time_of_year")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_recurrence_id"))
-  public static class Recurring extends BillingEvent {
+  public static class Recurring extends BillingEvent implements DatastoreAndSqlEntity {
 
     /**
      * The billing event recurs every year between {@link #eventTime} and this time on the
@@ -544,7 +545,7 @@ public abstract class BillingEvent extends ImmutableObject
         @javax.persistence.Index(columnList = "billingTime")
       })
   @AttributeOverride(name = "id", column = @Column(name = "billing_cancellation_id"))
-  public static class Cancellation extends BillingEvent {
+  public static class Cancellation extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The billing time of the charge that is being cancelled. */
     @Index
@@ -664,7 +665,7 @@ public abstract class BillingEvent extends ImmutableObject
   /** An event representing a modification of an existing one-time billing event. */
   @ReportedOn
   @Entity
-  public static class Modification extends BillingEvent {
+  public static class Modification extends BillingEvent implements DatastoreAndSqlEntity {
 
     /** The change in cost that should be applied to the original billing event. */
     Money cost;

core/src/main/java/google/registry/model/contact/ContactHistory.java

@@ -45,6 +45,7 @@ import javax.persistence.Id;
 @EntitySubclass
 @Access(AccessType.FIELD)
 public class ContactHistory extends HistoryEntry {
+
   // Store ContactBase instead of ContactResource so we don't pick up its @Id
   ContactBase contactBase;
 

core/src/main/java/google/registry/model/domain/DomainHistory.java

@@ -60,6 +60,7 @@ import javax.persistence.Table;
 @Access(AccessType.FIELD)
 @IdClass(DomainHistoryId.class)
 public class DomainHistory extends HistoryEntry {
+
   // Store DomainContent instead of DomainBase so we don't pick up its @Id
   DomainContent domainContent;
 

core/src/main/java/google/registry/model/index/EppResourceIndex.java

@@ -17,6 +17,7 @@ package google.registry.model.index;
 import static google.registry.util.TypeUtils.instantiate;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
@@ -25,11 +26,13 @@ import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.BackupGroupRoot;
 import google.registry.model.EppResource;
 import google.registry.model.annotations.ReportedOn;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
 
 /** An index that allows for quick enumeration of all EppResource entities (e.g. via map reduce). */
 @ReportedOn
 @Entity
-public class EppResourceIndex extends BackupGroupRoot {
+public class EppResourceIndex extends BackupGroupRoot implements DatastoreEntity {
 
   @Id
   String id;
@@ -74,4 +77,9 @@ public class EppResourceIndex extends BackupGroupRoot {
   public static <T extends EppResource> EppResourceIndex create(Key<T> resourceKey) {
     return create(EppResourceIndexBucket.getBucketKey(resourceKey), resourceKey);
   }
+
+  @Override
+  public ImmutableList<SqlEntity> toSqlEntities() {
+    return ImmutableList.of(); // not relevant in SQL
+  }
 }

core/src/main/java/google/registry/model/index/EppResourceIndexBucket.java

@@ -24,16 +24,23 @@ import com.googlecode.objectify.annotation.Id;
 import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.VirtualEntity;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
 
 /** A virtual entity to represent buckets to which EppResourceIndex objects are randomly added. */
 @Entity
 @VirtualEntity
-public class EppResourceIndexBucket extends ImmutableObject {
+public class EppResourceIndexBucket extends ImmutableObject implements DatastoreEntity {
 
   @SuppressWarnings("unused")
   @Id
   private long bucketId;
 
+  @Override
+  public ImmutableList<SqlEntity> toSqlEntities() {
+    return ImmutableList.of(); // not relevant in SQL
+  }
+
   /**
    * Deterministic function that returns a bucket id based on the resource's roid.
    * NB: At the moment, nothing depends on this being deterministic, so we have the ability to

core/src/main/java/google/registry/model/index/ForeignKeyIndex.java

@@ -43,6 +43,8 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
 import google.registry.util.NonFinalForTesting;
 import java.util.Map;
 import java.util.Optional;
@@ -61,28 +63,44 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
   /** The {@link ForeignKeyIndex} type for {@link ContactResource} entities. */
   @ReportedOn
   @Entity
-  public static class ForeignKeyContactIndex extends ForeignKeyIndex<ContactResource> {}
+  public static class ForeignKeyContactIndex extends ForeignKeyIndex<ContactResource>
+      implements DatastoreEntity {
+    @Override
+    public ImmutableList<SqlEntity> toSqlEntities() {
+      return ImmutableList.of(); // not relevant in SQL
+    }
+  }
 
   /** The {@link ForeignKeyIndex} type for {@link DomainBase} entities. */
   @ReportedOn
   @Entity
-  public static class ForeignKeyDomainIndex extends ForeignKeyIndex<DomainBase> {}
+  public static class ForeignKeyDomainIndex extends ForeignKeyIndex<DomainBase>
+      implements DatastoreEntity {
+    @Override
+    public ImmutableList<SqlEntity> toSqlEntities() {
+      return ImmutableList.of(); // not relevant in SQL
+    }
+  }
 
   /** The {@link ForeignKeyIndex} type for {@link HostResource} entities. */
   @ReportedOn
   @Entity
-  public static class ForeignKeyHostIndex extends ForeignKeyIndex<HostResource> {}
+  public static class ForeignKeyHostIndex extends ForeignKeyIndex<HostResource>
+      implements DatastoreEntity {
+    @Override
+    public ImmutableList<SqlEntity> toSqlEntities() {
+      return ImmutableList.of(); // not relevant in SQL
+    }
+  }
 
-  static final ImmutableMap<
-          Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
+  static final ImmutableMap<Class<? extends EppResource>, Class<? extends ForeignKeyIndex<?>>>
       RESOURCE_CLASS_TO_FKI_CLASS =
           ImmutableMap.of(
               ContactResource.class, ForeignKeyContactIndex.class,
               DomainBase.class, ForeignKeyDomainIndex.class,
               HostResource.class, ForeignKeyHostIndex.class);
 
-  @Id
-  String foreignKey;
+  @Id String foreignKey;
 
   /**
    * The deletion time of this {@link ForeignKeyIndex}.
@@ -90,8 +108,7 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    * <p>This will generally be equal to the deletion time of {@link #topReference}. However, in the
    * case of a {@link HostResource} that was renamed, this field will hold the time of the rename.
    */
-  @Index
-  DateTime deletionTime;
+  @Index DateTime deletionTime;
 
   /**
    * The referenced resource.

core/src/main/java/google/registry/model/poll/PollMessage.java

@@ -45,6 +45,7 @@ import google.registry.model.transfer.TransferResponse.ContactTransferResponse;
 import google.registry.model.transfer.TransferResponse.DomainTransferResponse;
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithLongVKey;
+import google.registry.schema.replay.DatastoreAndSqlEntity;
 import java.util.List;
 import java.util.Optional;
 import javax.persistence.AttributeOverride;
@@ -92,7 +93,7 @@ import org.joda.time.DateTime;
       @javax.persistence.Index(columnList = "eventTime")
     })
 public abstract class PollMessage extends ImmutableObject
-    implements Buildable, TransferServerApproveEntity {
+    implements Buildable, DatastoreAndSqlEntity, TransferServerApproveEntity {
 
   /** Entity id. */
   @Id

core/src/main/java/google/registry/model/registry/label/PremiumList.java

@@ -114,7 +114,7 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
   /** Virtual parent entity for premium list entry entities associated with a single revision. */
   @ReportedOn
   @Entity
-  public static class PremiumListRevision extends ImmutableObject {
+  public static class PremiumListRevision extends ImmutableObject implements DatastoreEntity {
 
     @Parent Key<PremiumList> parent;
 
@@ -171,6 +171,11 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
       }
       return revision;
     }
+
+    @Override
+    public ImmutableList<SqlEntity> toSqlEntities() {
+      return ImmutableList.of(); // not persisted in SQL
+    }
   }
 
   /**

core/src/main/java/google/registry/model/reporting/HistoryEntry.java

@@ -18,6 +18,7 @@ import static com.googlecode.objectify.Key.getKind;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
@@ -40,6 +41,8 @@ import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithStringVKey;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
 import java.util.Set;
 import javax.annotation.Nullable;
 import javax.persistence.Access;
@@ -59,7 +62,7 @@ import org.joda.time.DateTime;
 @MappedSuperclass
 @WithStringVKey // TODO(b/162229294): This should be resolved during the course of that bug
 @Access(AccessType.FIELD)
-public class HistoryEntry extends ImmutableObject implements Buildable {
+public class HistoryEntry extends ImmutableObject implements Buildable, DatastoreEntity, SqlEntity {
 
   /** Represents the type of history entry. */
   public enum Type {
@@ -307,6 +310,18 @@ public class HistoryEntry extends ImmutableObject implements Buildable {
     return resultEntity;
   }
 
+  // In SQL, save the child type
+  @Override
+  public ImmutableList<SqlEntity> toSqlEntities() {
+    return ImmutableList.of(toChildHistoryEntity());
+  }
+
+  // In Datastore, save as a HistoryEntry object regardless of this object's type
+  @Override
+  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
+    return ImmutableList.of(asHistoryEntry());
+  }
+
   /** A builder for {@link HistoryEntry} since it is immutable */
   public static class Builder<T extends HistoryEntry, B extends Builder<?, ?>>
       extends GenericBuilder<T, B> {

core/src/main/java/google/registry/persistence/transaction/TransactionEntity.java

@@ -14,6 +14,9 @@
 
 package google.registry.persistence.transaction;
 
+import com.google.common.collect.ImmutableList;
+import google.registry.schema.replay.DatastoreEntity;
+import google.registry.schema.replay.SqlEntity;
 import javax.persistence.Entity;
 import javax.persistence.GeneratedValue;
 import javax.persistence.GenerationType;
@@ -27,7 +30,7 @@ import javax.persistence.Table;
  */
 @Entity
 @Table(name = "Transaction")
-public class TransactionEntity {
+public class TransactionEntity implements SqlEntity {
 
   @Id
   @GeneratedValue(strategy = GenerationType.IDENTITY)
@@ -40,4 +43,9 @@ public class TransactionEntity {
   TransactionEntity(byte[] contents) {
     this.contents = contents;
   }
+
+  @Override
+  public ImmutableList<DatastoreEntity> toDatastoreEntities() {
+    return ImmutableList.of(); // not stored in Datastore per se
+  }
 }

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

@@ -21,6 +21,7 @@ import static org.joda.time.DateTimeZone.UTC;
 
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.common.CrossTldSingleton;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
 import google.registry.testing.AppEngineExtension;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.Test;
@@ -38,6 +39,7 @@ public class CreateAutoTimestampTest {
 
   /** Timestamped class. */
   @Entity(name = "CatTestEntity")
+  @EntityForTesting
   public static class TestObject extends CrossTldSingleton {
     CreateAutoTimestamp createTime = CreateAutoTimestamp.create(null);
   }

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

@@ -29,6 +29,7 @@ import com.google.common.collect.Iterables;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
 import google.registry.testing.AppEngineExtension;
 import google.registry.util.CidrAddressBlock;
 import java.util.ArrayDeque;
@@ -279,6 +280,7 @@ public class ImmutableObjectTest {
 
   /** Simple subclass of ImmutableObject. */
   @Entity
+  @EntityForTesting
   public static class ValueObject extends ImmutableObject {
     @Id
     long id;

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

@@ -21,6 +21,7 @@ import static org.joda.time.DateTimeZone.UTC;
 
 import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.common.CrossTldSingleton;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
 import google.registry.testing.AppEngineExtension;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.Test;
@@ -38,6 +39,7 @@ public class UpdateAutoTimestampTest {
 
   /** Timestamped class. */
   @Entity(name = "UatTestEntity")
+  @EntityForTesting
   public static class UpdateAutoTimestampTestObject extends CrossTldSingleton {
     UpdateAutoTimestamp updateTime = UpdateAutoTimestamp.create(null);
   }

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

@@ -45,6 +45,7 @@ import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.DatastoreHelper;
 import google.registry.testing.FakeClock;
@@ -69,13 +70,14 @@ public class OfyTest {
   @BeforeEach
   void beforeEach() {
     createTld("tld");
-    someObject = new HistoryEntry.Builder()
-        .setClientId("client id")
-        .setModificationTime(START_OF_TIME)
-        .setParent(persistActiveContact("parentContact"))
-        .setTrid(Trid.create("client", "server"))
-        .setXmlBytes("<xml></xml>".getBytes(UTF_8))
-        .build();
+    someObject =
+        new HistoryEntry.Builder()
+            .setClientId("client id")
+            .setModificationTime(START_OF_TIME)
+            .setParent(persistActiveContact("parentContact"))
+            .setTrid(Trid.create("client", "server"))
+            .setXmlBytes("<xml></xml>".getBytes(UTF_8))
+            .build();
     // This can't be initialized earlier because namespaces need the AppEngineRule to work.
   }
 
@@ -111,8 +113,7 @@ public class OfyTest {
         assertThrows(
             IllegalArgumentException.class,
             () ->
-                tm()
-                    .transact(
+                tm().transact(
                         () -> {
                           ofy().save().entity(someObject);
                           ofy().save().entity(someObject);
@@ -126,8 +127,7 @@ public class OfyTest {
         assertThrows(
             IllegalArgumentException.class,
             () ->
-                tm()
-                    .transact(
+                tm().transact(
                         () -> {
                           ofy().delete().entity(someObject);
                           ofy().delete().entity(someObject);
@@ -141,8 +141,7 @@ public class OfyTest {
         assertThrows(
             IllegalArgumentException.class,
             () ->
-                tm()
-                    .transact(
+                tm().transact(
                         () -> {
                           ofy().save().entity(someObject);
                           ofy().delete().entity(someObject);
@@ -156,8 +155,7 @@ public class OfyTest {
         assertThrows(
             IllegalArgumentException.class,
             () ->
-                tm()
-                    .transact(
+                tm().transact(
                         () -> {
                           ofy().delete().entity(someObject);
                           ofy().save().entity(someObject);
@@ -174,13 +172,12 @@ public class OfyTest {
 
   /** Simple entity class with lifecycle callbacks. */
   @com.googlecode.objectify.annotation.Entity
+  @EntityForTesting
   public static class LifecycleObject extends ImmutableObject {
 
-    @Parent
-    Key<?> parent = getCrossTldKey();
+    @Parent Key<?> parent = getCrossTldKey();
 
-    @Id
-    long id = 1;
+    @Id long id = 1;
 
     boolean onLoadCalled;
     boolean onSaveCalled;
@@ -218,20 +215,22 @@ public class OfyTest {
   /** Avoid regressions of b/21309102 where transaction time did not change on each retry. */
   @Test
   void testTransact_getsNewTimestampOnEachTry() {
-    tm().transact(new Runnable() {
+    tm().transact(
+            new Runnable() {
 
-      DateTime firstAttemptTime;
+              DateTime firstAttemptTime;
 
-      @Override
-      public void run() {
-        if (firstAttemptTime == null) {
-          // Sleep a bit to ensure that the next attempt is at a new millisecond.
-          firstAttemptTime = tm().getTransactionTime();
-          sleepUninterruptibly(10, MILLISECONDS);
-          throw new ConcurrentModificationException();
-        }
-        assertThat(tm().getTransactionTime()).isGreaterThan(firstAttemptTime);
-      }});
+              @Override
+              public void run() {
+                if (firstAttemptTime == null) {
+                  // Sleep a bit to ensure that the next attempt is at a new millisecond.
+                  firstAttemptTime = tm().getTransactionTime();
+                  sleepUninterruptibly(10, MILLISECONDS);
+                  throw new ConcurrentModificationException();
+                }
+                assertThat(tm().getTransactionTime()).isGreaterThan(firstAttemptTime);
+              }
+            });
   }
 
   @Test
@@ -319,17 +318,19 @@ public class OfyTest {
           }
         };
     // A commit logged work that throws on the first attempt to get its result.
-    CommitLoggedWork<Void> commitLoggedWork = new CommitLoggedWork<Void>(work, new SystemClock()) {
-      boolean firstCallToGetResult = true;
+    CommitLoggedWork<Void> commitLoggedWork =
+        new CommitLoggedWork<Void>(work, new SystemClock()) {
+          boolean firstCallToGetResult = true;
 
-      @Override
-      public Void getResult() {
-        if (firstCallToGetResult) {
-          firstCallToGetResult = false;
-          throw new DatastoreTimeoutException("");
-        }
-        return null;
-      }};
+          @Override
+          public Void getResult() {
+            if (firstCallToGetResult) {
+              firstCallToGetResult = false;
+              throw new DatastoreTimeoutException("");
+            }
+            return null;
+          }
+        };
     // Despite the DatastoreTimeoutException in the first call to getResult(), this should succeed
     // without retrying. If a retry is triggered, the test should fail due to the call to fail().
     ofy().transactCommitLoggedWork(commitLoggedWork);
@@ -381,8 +382,7 @@ public class OfyTest {
   void test_getBaseEntityClassFromEntityOrKey_subclassEntity() {
     DomainBase domain = DatastoreHelper.newDomainBase("test.tld");
     assertThat(getBaseEntityClassFromEntityOrKey(domain)).isEqualTo(DomainBase.class);
-    assertThat(getBaseEntityClassFromEntityOrKey(Key.create(domain)))
-        .isEqualTo(DomainBase.class);
+    assertThat(getBaseEntityClassFromEntityOrKey(Key.create(domain))).isEqualTo(DomainBase.class);
   }
 
   @Test

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

@@ -27,6 +27,7 @@ import com.googlecode.objectify.annotation.Entity;
 import google.registry.model.common.CrossTldSingleton;
 import google.registry.model.ofy.CommitLogManifest;
 import google.registry.model.ofy.Ofy;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectExtension;
@@ -42,6 +43,7 @@ public class CommitLogRevisionsTranslatorFactoryTest {
   private static final DateTime START_TIME = DateTime.parse("2000-01-01TZ");
 
   @Entity(name = "ClrtfTestEntity")
+  @EntityForTesting
   public static class TestObject extends CrossTldSingleton {
     ImmutableSortedMap<DateTime, Key<CommitLogManifest>> revisions = ImmutableSortedMap.of();
   }
@@ -73,11 +75,11 @@ public class CommitLogRevisionsTranslatorFactoryTest {
 
   @Test
   void testSave_doesNotMutateOriginalResource() {
-     TestObject object = new TestObject();
-     save(object);
-     assertThat(object.revisions).isEmpty();
-     assertThat(reload().revisions).isNotEmpty();
-   }
+    TestObject object = new TestObject();
+    save(object);
+    assertThat(object.revisions).isEmpty();
+    assertThat(reload().revisions).isNotEmpty();
+  }
 
   @Test
   void testSave_translatorAddsKeyToCommitLogToField() {
@@ -149,8 +151,10 @@ public class CommitLogRevisionsTranslatorFactoryTest {
     com.google.appengine.api.datastore.Entity entity =
         tm().transactNewReadOnly(() -> ofy().save().toEntity(reload()));
     assertThat(entity.getProperties().keySet()).containsExactly("revisions.key", "revisions.value");
-    assertThat(entity.getProperties()).containsEntry(
-        "revisions.key", ImmutableList.of(START_TIME.toDate(), START_TIME.plusDays(1).toDate()));
+    assertThat(entity.getProperties())
+        .containsEntry(
+            "revisions.key",
+            ImmutableList.of(START_TIME.toDate(), START_TIME.plusDays(1).toDate()));
     assertThat(entity.getProperty("revisions.value")).isInstanceOf(List.class);
     assertThat(((List<Object>) entity.getProperty("revisions.value")).get(0))
         .isInstanceOf(com.google.appengine.api.datastore.Key.class);

core/src/test/java/google/registry/schema/replay/EntityTest.java

@@ -41,24 +41,32 @@ public class EntityTest {
         new ClassGraph().enableAnnotationInfo().whitelistPackages("google.registry").scan()) {
       // All javax.persistence entities must implement SqlEntity and vice versa
       ImmutableSet<String> javaxPersistenceClasses =
-          getClassNames(
-              scanResult.getClassesWithAnnotation(javax.persistence.Entity.class.getName()));
+          getAllClassesWithAnnotation(scanResult, javax.persistence.Entity.class.getName());
       ImmutableSet<String> sqlEntityClasses =
           getClassNames(scanResult.getClassesImplementing(SqlEntity.class.getName()));
-      assertThat(javaxPersistenceClasses).isEqualTo(sqlEntityClasses);
+      assertThat(sqlEntityClasses).containsExactlyElementsIn(javaxPersistenceClasses);
 
-      // All com.googlecode.objectify.annotation.Entity classes must implement DatastoreEntity and
-      // vice versa
+      // All com.googlecode.objectify entities must implement DatastoreEntity and vice versa
       ImmutableSet<String> objectifyClasses =
-          getClassNames(
-              scanResult.getClassesWithAnnotation(
-                  com.googlecode.objectify.annotation.Entity.class.getName()));
+          getAllClassesWithAnnotation(
+              scanResult, com.googlecode.objectify.annotation.Entity.class.getName());
       ImmutableSet<String> datastoreEntityClasses =
           getClassNames(scanResult.getClassesImplementing(DatastoreEntity.class.getName()));
-      assertThat(objectifyClasses).isEqualTo(datastoreEntityClasses);
+      assertThat(datastoreEntityClasses).containsExactlyElementsIn(objectifyClasses);
     }
   }
 
+  private ImmutableSet<String> getAllClassesWithAnnotation(
+      ScanResult scanResult, String annotation) {
+    ImmutableSet.Builder<String> result = new ImmutableSet.Builder<>();
+    ClassInfoList classesWithAnnotation = scanResult.getClassesWithAnnotation(annotation);
+    result.addAll(getClassNames(classesWithAnnotation));
+    classesWithAnnotation.stream()
+        .map(ClassInfo::getSubclasses)
+        .forEach(classInfoList -> result.addAll(getClassNames(classInfoList)));
+    return result.build();
+  }
+
   private ImmutableSet<String> getClassNames(ClassInfoList classInfoList) {
     return classInfoList.stream()
         .filter(ClassInfo::isStandardClass)

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

@@ -23,11 +23,11 @@ import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.VirtualEntity;
 import google.registry.model.common.EntityGroupRoot;
+import google.registry.schema.replay.EntityTest.EntityForTesting;
 
-/**
- * A test model object that can be persisted in any entity group.
- */
+/** A test model object that can be persisted in any entity group. */
 @Entity
+@EntityForTesting
 public class TestObject extends ImmutableObject {
 
   @Parent
@@ -65,6 +65,7 @@ public class TestObject extends ImmutableObject {
   /** A test @VirtualEntity model object, which should not be persisted. */
   @Entity
   @VirtualEntity
+  @EntityForTesting
   public static class TestVirtualObject extends ImmutableObject {
 
     @Id