Allow addition of extra entity classes for VKey conversion (#877)

* Allow addition of extra entity classes for VKey conversion

This allows us to create VKeys from Keys for test objects that may not
be part of the original codebase.

This isn't used anywhere directly yet but it will be useful in the
future when testing the replay of SQL transactions.

core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java

@@ -16,15 +16,16 @@ package google.registry.model.translators;
 
 import static com.google.common.base.Functions.identity;
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static google.registry.model.EntityClasses.ALL_CLASSES;
 
 import com.google.appengine.api.datastore.Key;
-import com.google.common.collect.ImmutableMap;
+import com.google.common.annotations.VisibleForTesting;
 import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.persistence.VKey;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -39,10 +40,10 @@ public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory<VKey,
   // name, which is all the datastore key gives us.
   // Note that entities annotated with @EntitySubclass are removed because they share the same
   // kind of the key with their parent class.
-  private static final ImmutableMap<String, Class<?>> CLASS_REGISTRY =
+  private static final Map<String, Class<?>> CLASS_REGISTRY =
       ALL_CLASSES.stream()
           .filter(clazz -> !clazz.isAnnotationPresent(EntitySubclass.class))
-          .collect(toImmutableMap(com.googlecode.objectify.Key::getKind, identity()));
+          .collect(Collectors.toMap(com.googlecode.objectify.Key::getKind, identity()));
 
   public VKeyTranslatorFactory() {
     super(VKey.class);
@@ -59,6 +60,7 @@ public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory<VKey,
 
   /** Create a VKey from an objectify Key. */
   @Nullable
+  @SuppressWarnings("unchecked")
   public static <T> VKey<T> createVKey(@Nullable com.googlecode.objectify.Key<T> key) {
     if (key == null) {
       return null;
@@ -95,6 +97,11 @@ public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory<VKey,
     return createVKey(com.googlecode.objectify.Key.create(urlSafe));
   }
 
+  @VisibleForTesting
+  public static void addTestEntityClass(Class<?> clazz) {
+    CLASS_REGISTRY.put(com.googlecode.objectify.Key.getKind(clazz), clazz);
+  }
+
   @Override
   public SimpleTranslator<VKey, Key> createTranslator() {
     return new SimpleTranslator<VKey, Key>() {

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

@@ -27,16 +27,24 @@ import google.registry.model.ofy.CommitLogCheckpointRoot;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.TestObject;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 public class VKeyTranslatorFactoryTest {
 
   @RegisterExtension
-  public final AppEngineExtension appEngine = AppEngineExtension.builder().withDatastore().build();
+  public final AppEngineExtension appEngine =
+      AppEngineExtension.builder().withDatastore().withOfyTestEntities(TestObject.class).build();
 
   VKeyTranslatorFactoryTest() {}
 
+  @BeforeAll
+  static void beforeAll() {
+    VKeyTranslatorFactory.addTestEntityClass(TestObject.class);
+  }
+
   @Test
   void testEntityWithFlatKey() {
     // Creating an objectify key instead of a datastore key as this should get a correctly formatted
@@ -88,4 +96,11 @@ public class VKeyTranslatorFactoryTest {
     assertThat(vkey.getOfyKey()).isEqualTo(key);
     assertThat(vkey.getSqlKey()).isEqualTo("ROID-1");
   }
+
+  @Test
+  void testExtraEntityClass() {
+    TestObject testObject = TestObject.create("id", "field");
+    Key<TestObject> key = Key.create(testObject);
+    assertThat(VKeyTranslatorFactory.createVKey(key).getSqlKey()).isEqualTo("id");
+  }
 }

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

@@ -23,6 +23,7 @@ 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.persistence.VKey;
 import google.registry.schema.replay.EntityTest.EntityForTesting;
 
 /** A test model object that can be persisted in any entity group. */
@@ -30,11 +31,9 @@ import google.registry.schema.replay.EntityTest.EntityForTesting;
 @EntityForTesting
 public class TestObject extends ImmutableObject {
 
-  @Parent
-  Key<EntityGroupRoot> parent;
+  @Parent Key<EntityGroupRoot> parent;
 
-  @Id
-  String id;
+  @Id String id;
 
   String field;
 
@@ -46,6 +45,10 @@ public class TestObject extends ImmutableObject {
     return field;
   }
 
+  public static VKey<TestObject> createVKey(Key<TestObject> key) {
+    return VKey.create(TestObject.class, key.getName(), key);
+  }
+
   public static TestObject create(String id) {
     return create(id, null);
   }
@@ -68,8 +71,7 @@ public class TestObject extends ImmutableObject {
   @EntityForTesting
   public static class TestVirtualObject extends ImmutableObject {
 
-    @Id
-    String id;
+    @Id String id;
 
     /**
      * Expose a factory method for testing saves of virtual entities; in real life this would never