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

docs/app-engine-architecture.md

@@ -1,6 +1,7 @@
 # App Engine architecture
 
-This document contains information on the overall architecture of the Domain Registry project as it is implemented in App Engine.
+This document contains information on the overall architecture of the Domain
+Registry project as it is implemented in App Engine.
 
 ## Modules
 

docs/developing.md

@@ -1,3 +1,4 @@
 # Developing
 
-Advice on how to do development on the Domain Registry codebase (including how to set up an IDE environment and run tests).
+Advice on how to do development on the Domain Registry codebase (including how
+to set up an IDE environment and run tests).

docs/install.md

@@ -1,3 +1,4 @@
 # Installation
 
-Information on how to download and install the Domain Registry project and get a working running instance.
+Information on how to download and install the Domain Registry project and get a
+working running instance.

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(

javatests/google/registry/tools/CreateTldCommandTest.java

@@ -241,6 +241,16 @@ public class CreateTldCommandTest extends CommandTestCase<CreateTldCommand> {
     runCommandForced("--initial_tld_state=INVALID_STATE", "--roid_suffix=Q9JYB4C", "xn--q9jyb4c");
   }
 
+  @Test
+  public void testFailure_bothTldStateFlags() throws Exception {
+    thrown.expect(IllegalArgumentException.class);
+    DateTime now = DateTime.now(UTC);
+    runCommandForced(
+        String.format("--tld_state_transitions=%s=PREDELEGATION,%s=SUNRISE", now, now.plus(1)),
+        "--initial_tld_state=GENERAL_AVAILABILITY",
+        "xn--q9jyb4c");
+    }
+
   @Test
   public void testFailure_negativeInitialRenewBillingCost() throws Exception {
     thrown.expect(IllegalArgumentException.class);

javatests/google/registry/tools/UpdateTldCommandTest.java

@@ -94,6 +94,17 @@ public class UpdateTldCommandTest extends CommandTestCase<UpdateTldCommand> {
     assertThat(registry.getTldState(END_OF_TIME)).isEqualTo(TldState.GENERAL_AVAILABILITY);
   }
 
+  @Test
+  public void testSuccess_setTldState() throws Exception {
+    Registry registry = persistResource(
+        Registry.get("xn--q9jyb4c").asBuilder()
+            .setTldStateTransitions(ImmutableSortedMap.of(START_OF_TIME, TldState.PREDELEGATION))
+            .build());
+    runCommandForced("--set_current_tld_state=SUNRISE", "xn--q9jyb4c");
+    registry = Registry.get("xn--q9jyb4c");
+    assertThat(registry.getTldState(now.plusDays(1))).isEqualTo(TldState.SUNRISE);
+  }
+
   @Test
   public void testSuccess_renewBillingCostTransitions() throws Exception {
     DateTime later = now.plusMonths(1);
@@ -457,6 +468,65 @@ public class UpdateTldCommandTest extends CommandTestCase<UpdateTldCommand> {
         "xn--q9jyb4c");
   }
 
+  @Test
+  public void testFailure_bothTldStateFlags() throws Exception {
+    thrown.expect(
+        IllegalArgumentException.class,
+        "Don't pass both --set_current_tld_state and --tld_state_transitions");
+    runCommandForced(
+        String.format("--tld_state_transitions=%s=PREDELEGATION,%s=SUNRISE", now, now.plusDays(1)),
+        "--set_current_tld_state=GENERAL_AVAILABILITY",
+        "xn--q9jyb4c");
+    }
+
+  @Test
+  public void testFailure_setCurrentTldState_outOfOrder() throws Exception {
+    thrown.expect(
+        IllegalArgumentException.class, "The TLD states are chronologically out of order");
+    persistResource(
+        Registry.get("xn--q9jyb4c").asBuilder()
+            .setTldStateTransitions(
+                ImmutableSortedMap.of(
+                    START_OF_TIME, TldState.PREDELEGATION,
+                    now.minusMonths(1), TldState.GENERAL_AVAILABILITY))
+            .build());
+    runCommandForced("--set_current_tld_state=SUNRISE", "xn--q9jyb4c");
+  }
+
+  @Test
+  public void testFailure_setCurrentTldState_laterTransitionScheduled() throws Exception {
+    thrown.expect(
+        IllegalArgumentException.class,
+        " when there is a later transition already scheduled");
+    persistResource(
+        Registry.get("xn--q9jyb4c").asBuilder()
+            .setTldStateTransitions(
+                ImmutableSortedMap.of(
+                    START_OF_TIME, TldState.PREDELEGATION,
+                    now.plusMonths(1), TldState.GENERAL_AVAILABILITY))
+            .build());
+    runCommandForced("--set_current_tld_state=SUNRISE", "xn--q9jyb4c");
+  }
+
+  @Test
+  public void testFailure_setCurrentTldState_inProduction() throws Exception {
+    thrown.expect(
+        IllegalArgumentException.class,
+        "--set_current_tld_state is not safe to use in production.");
+    persistResource(
+        Registry.get("xn--q9jyb4c").asBuilder()
+            .setTldStateTransitions(
+                ImmutableSortedMap.of(
+                    START_OF_TIME, TldState.PREDELEGATION,
+                    now.minusMonths(1), TldState.GENERAL_AVAILABILITY))
+            .build());
+    runCommandInEnvironment(
+        RegistryToolEnvironment.PRODUCTION,
+        "--set_current_tld_state=SUNRISE",
+        "xn--q9jyb4c",
+        "--force");
+  }
+
   @Test
   public void testFailure_invalidRenewBillingCost() throws Exception {
     thrown.expect(ParameterException.class);