Add --set_current_tld_state to UpdateTldCommand

This feature would have been useful earlier when I was changing the TLD state on a sandbox TLD on-the-fly for testing purposes.

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

java/google/registry/tools/CreateOrUpdateTldCommand.java

@@ -36,6 +36,7 @@ import google.registry.tools.params.OptionalStringParameter;
 import google.registry.tools.params.TransitionListParameter.BillingCostTransitions;
 import google.registry.tools.params.TransitionListParameter.TldStateTransitions;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.annotation.Nullable;
 import org.joda.money.Money;
@@ -214,6 +215,8 @@ abstract class CreateOrUpdateTldCommand extends MutatingCommand {
 
   abstract ImmutableSet<String> getReservedLists(Registry oldRegistry);
 
+  abstract Optional<Map.Entry<DateTime, TldState>> getTldStateTransitionToAdd();
+
   /** Subclasses can override this to set their own properties. */
   void setCommandSpecificProperties(@SuppressWarnings("unused") Registry.Builder builder) {}
 
@@ -260,8 +263,24 @@ abstract class CreateOrUpdateTldCommand extends MutatingCommand {
         builder.setDnsPaused(!dns);
       }
 
+      Optional<Map.Entry<DateTime, TldState>> tldStateTransitionToAdd =
+          getTldStateTransitionToAdd();
       if (!tldStateTransitions.isEmpty()) {
         builder.setTldStateTransitions(tldStateTransitions);
+      } else if (tldStateTransitionToAdd.isPresent()) {
+        ImmutableSortedMap.Builder<DateTime, TldState> newTldStateTransitions =
+            ImmutableSortedMap.naturalOrder();
+        if (oldRegistry != null) {
+          checkArgument(
+              oldRegistry.getTldStateTransitions().lastKey().isBefore(
+                  tldStateTransitionToAdd.get().getKey()),
+              "Cannot add %s at %s when there is a later transition already scheduled",
+              tldStateTransitionToAdd.get().getValue(),
+              tldStateTransitionToAdd.get().getKey());
+          newTldStateTransitions.putAll(oldRegistry.getTldStateTransitions());
+        }
+        builder.setTldStateTransitions(
+            newTldStateTransitions.put(getTldStateTransitionToAdd().get()).build());
       }
 
       if (!renewBillingCostTransitions.isEmpty()) {

java/google/registry/tools/CreateTldCommand.java

@@ -22,14 +22,18 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import com.google.common.base.Optional;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Maps;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
+import java.util.Map;
 import javax.annotation.Nullable;
 import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
+import org.joda.time.DateTime;
 
 /** Command to create a TLD. */
 @Parameters(separators = " =", commandDescription = "Create new TLD(s)")
@@ -52,9 +56,6 @@ class CreateTldCommand extends CreateOrUpdateTldCommand {
   protected void initTldCommand() throws Exception {
     checkArgument(initialTldState == null || tldStateTransitions.isEmpty(),
         "Don't pass both --initial_tld_state and --tld_state_transitions");
-    if (initialTldState != null) {
-      tldStateTransitions = ImmutableSortedMap.of(START_OF_TIME, initialTldState);
-    }
     checkArgument(initialRenewBillingCost == null || renewBillingCostTransitions.isEmpty(),
         "Don't pass both --initial_renew_billing_cost and --renew_billing_cost_transitions");
     if (initialRenewBillingCost != null) {
@@ -78,7 +79,7 @@ class CreateTldCommand extends CreateOrUpdateTldCommand {
 
     // If this is a non-default currency and the user hasn't specified an EAP fee schedule, set the
     // EAP fee schedule to a matching currency.
-    if (currency != Registry.DEFAULT_CURRENCY && eapFeeSchedule.isEmpty()) {
+    if (!currency.equals(Registry.DEFAULT_CURRENCY) && eapFeeSchedule.isEmpty()) {
       builder.setEapFeeSchedule(ImmutableSortedMap.of(START_OF_TIME, Money.zero(currency)));
     }
   }
@@ -103,4 +104,11 @@ class CreateTldCommand extends CreateOrUpdateTldCommand {
   ImmutableSet<String> getReservedLists(Registry oldRegistry) {
     return ImmutableSet.copyOf(nullToEmpty(reservedListNames));
   }
+
+  @Override
+  Optional<Map.Entry<DateTime, TldState>> getTldStateTransitionToAdd() {
+    return initialTldState != null
+        ? Optional.of(Maps.immutableEntry(START_OF_TIME, initialTldState))
+        : Optional.<Map.Entry<DateTime, TldState>>absent();
+  }
 }

java/google/registry/tools/UpdateTldCommand.java

@@ -24,12 +24,20 @@ import static google.registry.util.CollectionUtils.nullToEmpty;
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.base.Function;
+import com.google.common.base.Optional;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Maps;
 import com.googlecode.objectify.Key;
+import google.registry.config.RegistryEnvironment;
 import google.registry.model.registry.Registry;
+import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.label.ReservedList;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
+
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.annotation.Nullable;
 
@@ -72,6 +80,13 @@ class UpdateTldCommand extends CreateOrUpdateTldCommand {
       description = "A comma-separated list of allowed nameservers to be removed from the TLD")
   List<String> allowedNameserversRemove;
 
+  @Nullable
+  @Parameter(
+      names = "--set_current_tld_state",
+      description = "Set the current TLD state. Specifically, adds a TLD transition at the "
+          + "current time for the specified state.")
+  TldState setCurrentTldState;
+
   @Override
   Registry getOldRegistry(String tld) {
     return Registry.get(assertTldExists(tld));
@@ -115,13 +130,28 @@ class UpdateTldCommand extends CreateOrUpdateTldCommand {
         reservedListsRemove);
   }
 
+  @Override
+  Optional<Map.Entry<DateTime, TldState>> getTldStateTransitionToAdd() {
+    return setCurrentTldState != null
+        ? Optional.of(Maps.immutableEntry(DateTime.now(DateTimeZone.UTC), setCurrentTldState))
+        : Optional.<Map.Entry<DateTime, TldState>>absent();
+  }
+
   @Override
   protected void initTldCommand() throws Exception {
+    // Due to per-instance caching on Registry, different instances can end up in different TLD
+    // states at the same time, so --set_current_tld_state should never be used in production.
+    checkArgument(
+        !RegistryEnvironment.get().equals(RegistryEnvironment.PRODUCTION)
+            || setCurrentTldState == null,
+        "--set_current_tld_state is not safe to use in production.");
     checkConflicts("reserved_lists", reservedListNames, reservedListsAdd, reservedListsRemove);
     checkConflicts(
         "allowed_registrants", allowedRegistrants, allowedRegistrantsAdd, allowedRegistrantsRemove);
     checkConflicts(
         "allowed_nameservers", allowedNameservers, allowedNameserversAdd, allowedNameserversRemove);
+    checkArgument(setCurrentTldState == null || tldStateTransitions.isEmpty(),
+        "Don't pass both --set_current_tld_state and --tld_state_transitions");
   }
 
   private static ImmutableSet<String> formUpdatedList(