diff --git a/core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java b/core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java index fd529f2b8..99615f879 100644 --- a/core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java +++ b/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> CLASS_REGISTRY = + private static final Map> 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 createVKey(@Nullable com.googlecode.objectify.Key key) { if (key == null) { return null; @@ -95,6 +97,11 @@ public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory clazz) { + CLASS_REGISTRY.put(com.googlecode.objectify.Key.getKind(clazz), clazz); + } + @Override public SimpleTranslator createTranslator() { return new SimpleTranslator() { diff --git a/core/src/test/java/google/registry/model/translators/VKeyTranslatorFactoryTest.java b/core/src/test/java/google/registry/model/translators/VKeyTranslatorFactoryTest.java index 4f9df0c5d..15cdaea1e 100644 --- a/core/src/test/java/google/registry/model/translators/VKeyTranslatorFactoryTest.java +++ b/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 key = Key.create(testObject); + assertThat(VKeyTranslatorFactory.createVKey(key).getSqlKey()).isEqualTo("id"); + } } diff --git a/core/src/test/java/google/registry/testing/TestObject.java b/core/src/test/java/google/registry/testing/TestObject.java index 7392ccf9a..15b2e368f 100644 --- a/core/src/test/java/google/registry/testing/TestObject.java +++ b/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 parent; + @Parent Key parent; - @Id - String id; + @Id String id; String field; @@ -46,6 +45,10 @@ public class TestObject extends ImmutableObject { return field; } + public static VKey createVKey(Key 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