Perform minor refactors on premium list code

Principally, this moves a load method into DatastoreHelper that is now
only used by tests.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=148649087

java/google/registry/model/registry/label/PremiumList.java

@@ -39,7 +39,6 @@ import com.googlecode.objectify.annotation.Cache;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Parent;
-import com.googlecode.objectify.cmd.Query;
 import google.registry.model.Buildable;
 import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.ReportedOn;
@@ -287,10 +286,6 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
     return Objects.equals(registry.getPremiumList(), key);
   }
 
-  Query<PremiumListEntry> queryEntriesForCurrentRevision() {
-    return ofy().load().type(PremiumListEntry.class).ancestor(revisionKey);
-  }
-
   @Override
   public Builder asBuilder() {
     return new Builder(clone(this));

java/google/registry/model/registry/label/PremiumListUtils.java

@@ -14,7 +14,6 @@
 
 package google.registry.model.registry.label;
 
-import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.Iterables.partition;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
@@ -37,7 +36,6 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.PremiumList.PremiumListEntry;
 import google.registry.model.registry.label.PremiumList.PremiumListRevision;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ExecutionException;
 import org.joda.money.Money;
@@ -107,9 +105,7 @@ public final class PremiumListUtils {
         PremiumListRevision.create(premiumList, premiumListEntries.keySet());
     final Key<PremiumListRevision> newRevisionKey = Key.create(newRevision);
     ImmutableSet<PremiumListEntry> parentedEntries =
-        parentPremiumListEntriesOnRevision(
+        parentPremiumListEntriesOnRevision(premiumListEntries.values(), newRevisionKey);
-            firstNonNull(premiumListEntries.values(), ImmutableSet.<PremiumListEntry>of()),
-            newRevisionKey);
 
     // Save the new child entities in a series of transactions.
     for (final List<PremiumListEntry> batch :
@@ -159,6 +155,7 @@ public final class PremiumListUtils {
   }
 
   /** Re-parents the given {@link PremiumListEntry}s on the given {@link PremiumListRevision}. */
+  @VisibleForTesting
   public static ImmutableSet<PremiumListEntry> parentPremiumListEntriesOnRevision(
       Iterable<PremiumListEntry> entries, final Key<PremiumListRevision> revisionKey) {
     return FluentIterable.from(entries)
@@ -188,7 +185,7 @@ public final class PremiumListUtils {
       return;
     }
     for (final List<Key<PremiumListEntry>> batch : partition(
-        premiumList.queryEntriesForCurrentRevision().keys(),
+        ofy().load().type(PremiumListEntry.class).ancestor(premiumList.revisionKey).keys(),
         TRANSACTION_BATCH_SIZE)) {
       ofy().transactNew(new VoidWork() {
         @Override
@@ -208,28 +205,5 @@ public final class PremiumListUtils {
     return ofy().load().key(Key.create(getCrossTldKey(), PremiumList.class, name)).now() != null;
   }
 
-  /**
-   * Loads and returns the entire premium list map.
-   *
-   * <p>This load operation is quite expensive for large premium lists because each premium list
-   * entry is a separate Datastore entity, and loading them this way bypasses the in-memory caches.
-   * Do not use this method if all you need to do is check the price of a small number of labels!
-   */
-  @VisibleForTesting
-  public static Map<String, PremiumListEntry> loadPremiumListEntries(PremiumList premiumList) {
-    try {
-      ImmutableMap.Builder<String, PremiumListEntry> entriesMap = new ImmutableMap.Builder<>();
-      if (premiumList.getRevisionKey() != null) {
-        for (PremiumListEntry entry : premiumList.queryEntriesForCurrentRevision()) {
-          entriesMap.put(entry.getLabel(), entry);
-        }
-      }
-      return entriesMap.build();
-    } catch (Exception e) {
-      throw new RuntimeException(
-          "Could not retrieve entries for premium list " + premiumList.getName(), e);
-    }
-  }
-
   private PremiumListUtils() {}
 }

javatests/google/registry/model/registry/label/PremiumListUtilsTest.java

@@ -20,13 +20,14 @@ import static google.registry.model.registry.label.PremiumList.cachePremiumListE
 import static google.registry.model.registry.label.PremiumListUtils.deletePremiumList;
 import static google.registry.model.registry.label.PremiumListUtils.doesPremiumListExist;
 import static google.registry.model.registry.label.PremiumListUtils.getPremiumPrice;
-import static google.registry.model.registry.label.PremiumListUtils.loadPremiumListEntries;
 import static google.registry.model.registry.label.PremiumListUtils.savePremiumListAndEntries;
 import static google.registry.testing.DatastoreHelper.createTld;
+import static google.registry.testing.DatastoreHelper.loadPremiumListEntries;
 import static google.registry.testing.DatastoreHelper.persistPremiumList;
 import static google.registry.testing.DatastoreHelper.persistResource;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.googlecode.objectify.Key;
 import google.registry.model.pricing.StaticPremiumListPricingEngine;
 import google.registry.model.registry.Registry;
@@ -187,7 +188,8 @@ public class PremiumListUtilsTest {
     persistResource(Registry.get("tld").asBuilder().setPremiumList(pl).build());
     assertThat(getPremiumPrice("lol", Registry.get("tld"))).hasValue(Money.parse("USD 999"));
     assertThat(getPremiumPrice("lol ", Registry.get("tld"))).isAbsent();
-    Map<String, PremiumListEntry> entries = loadPremiumListEntries(PremiumList.get("tld2").get());
+    ImmutableMap<String, PremiumListEntry> entries =
+        loadPremiumListEntries(PremiumList.get("tld2").get());
     assertThat(entries.keySet()).containsExactly("lol");
     assertThat(entries).doesNotContainKey("lol ");
     PremiumListEntry entry = entries.get("lol");
@@ -201,7 +203,7 @@ public class PremiumListUtilsTest {
     PremiumList pl =
         savePremiumListAndEntries(
             new PremiumList.Builder().setName("pl").build(), ImmutableList.of("test,USD 1"));
-    Map<String, PremiumListEntry> entries = loadPremiumListEntries(pl);
+    ImmutableMap<String, PremiumListEntry> entries = loadPremiumListEntries(pl);
     assertThat(entries.keySet()).containsExactly("test");
     assertThat(loadPremiumListEntries(PremiumList.get("pl").get())).isEqualTo(entries);
     // Save again with no changes, and clear the cache to force a re-load from Datastore.

javatests/google/registry/testing/DatastoreHelper.java

@@ -987,5 +987,23 @@ public class DatastoreHelper {
       }});
   }
 
+  /** Returns the entire map of {@link PremiumListEntry}s for the given {@link PremiumList}. */
+  public static ImmutableMap<String, PremiumListEntry> loadPremiumListEntries(
+      PremiumList premiumList) {
+    try {
+      ImmutableMap.Builder<String, PremiumListEntry> entriesMap = new ImmutableMap.Builder<>();
+      if (premiumList.getRevisionKey() != null) {
+        for (PremiumListEntry entry :
+            ofy().load().type(PremiumListEntry.class).ancestor(premiumList.getRevisionKey())) {
+          entriesMap.put(entry.getLabel(), entry);
+        }
+      }
+      return entriesMap.build();
+    } catch (Exception e) {
+      throw new RuntimeException(
+          "Could not retrieve entries for premium list " + premiumList.getName(), e);
+    }
+  }
+
   private DatastoreHelper() {}
 }

javatests/google/registry/tools/DeletePremiumListCommandTest.java

@@ -17,8 +17,8 @@ package google.registry.tools;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.model.ofy.ObjectifyService.ofy;
-import static google.registry.model.registry.label.PremiumListUtils.loadPremiumListEntries;
 import static google.registry.testing.DatastoreHelper.createTld;
+import static google.registry.testing.DatastoreHelper.loadPremiumListEntries;
 import static google.registry.testing.DatastoreHelper.persistPremiumList;
 import static google.registry.testing.DatastoreHelper.persistResource;
 

javatests/google/registry/tools/server/CreatePremiumListActionTest.java

@@ -17,8 +17,8 @@ package google.registry.tools.server;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.registry.label.PremiumListUtils.deletePremiumList;
 import static google.registry.model.registry.label.PremiumListUtils.getPremiumPrice;
-import static google.registry.model.registry.label.PremiumListUtils.loadPremiumListEntries;
 import static google.registry.testing.DatastoreHelper.createTlds;
+import static google.registry.testing.DatastoreHelper.loadPremiumListEntries;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 
 import google.registry.model.registry.Registry;

javatests/google/registry/tools/server/UpdatePremiumListActionTest.java

@@ -16,8 +16,8 @@ package google.registry.tools.server;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.registry.label.PremiumListUtils.getPremiumPrice;
-import static google.registry.model.registry.label.PremiumListUtils.loadPremiumListEntries;
 import static google.registry.testing.DatastoreHelper.createTlds;
+import static google.registry.testing.DatastoreHelper.loadPremiumListEntries;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
 
 import google.registry.model.registry.Registry;