From bca2169c6815f3418c0c5fdb6cdec5ea559998e1 Mon Sep 17 00:00:00 2001
From: ctingue <ctingue@google.com>
Date: Mon, 9 Jan 2017 08:28:29 -0800
Subject: [PATCH] Add global cursor functionality to UpdateCursorsCommand

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=143965268
---
 .../registry/tools/UpdateCursorsCommand.java  | 22 +++---
 .../tools/UpdateCursorsCommandTest.java       | 77 +++++++++++++++++--
 2 files changed, 84 insertions(+), 15 deletions(-)

diff --git a/java/google/registry/tools/UpdateCursorsCommand.java b/java/google/registry/tools/UpdateCursorsCommand.java
index 3d30b2434..80391396c 100644
--- a/java/google/registry/tools/UpdateCursorsCommand.java
+++ b/java/google/registry/tools/UpdateCursorsCommand.java
@@ -15,6 +15,7 @@
 package google.registry.tools;
 
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.util.CollectionUtils.isNullOrEmpty;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
@@ -29,9 +30,7 @@ import org.joda.time.DateTime;
 @Parameters(separators = " =", commandDescription = "Modifies cursor timestamps used by LRC tasks")
 final class UpdateCursorsCommand extends MutatingCommand {
 
-  @Parameter(
-      description = "TLDs on which to operate.",
-      required = true)
+  @Parameter(description = "TLDs on which to operate. Omit for global cursors.")
   private List<String> tlds;
 
   @Parameter(
@@ -49,12 +48,17 @@ final class UpdateCursorsCommand extends MutatingCommand {
 
   @Override
   protected void init() throws Exception {
-    for (String tld : tlds) {
-      Registry registry = Registry.get(tld);
-      Cursor cursor = ofy().load().key(Cursor.createKey(cursorType, registry)).now();
-      stageEntityChange(
-          cursor,
-          Cursor.create(cursorType, newTimestamp, registry));
+    if (isNullOrEmpty(tlds)) {
+      Cursor cursor = ofy().load().key(Cursor.createGlobalKey(cursorType)).now();
+      stageEntityChange(cursor, Cursor.createGlobal(cursorType, newTimestamp));
+    } else {
+      for (String tld : tlds) {
+        Registry registry = Registry.get(tld);
+        Cursor cursor = ofy().load().key(Cursor.createKey(cursorType, registry)).now();
+        stageEntityChange(
+            cursor,
+            Cursor.create(cursorType, newTimestamp, registry));
+      }
     }
   }
 }
diff --git a/javatests/google/registry/tools/UpdateCursorsCommandTest.java b/javatests/google/registry/tools/UpdateCursorsCommandTest.java
index e3c21fa37..3a604b83d 100644
--- a/javatests/google/registry/tools/UpdateCursorsCommandTest.java
+++ b/javatests/google/registry/tools/UpdateCursorsCommandTest.java
@@ -19,9 +19,11 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.persistResource;
 
+import com.beust.jcommander.ParameterException;
 import google.registry.model.common.Cursor;
 import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.registry.Registry;
+import google.registry.model.registry.Registry.RegistryNotFoundException;
 import org.joda.time.DateTime;
 import org.junit.Before;
 import org.junit.Test;
@@ -43,15 +45,78 @@ public class UpdateCursorsCommandTest extends CommandTestCase<UpdateCursorsComma
         .isEqualTo(DateTime.parse("1984-12-18TZ"));
   }
 
-  @Test
-  public void testUpdateCursors_oldValueIsAbsent() throws Exception {
-    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry)).now()).isNull();
-    doUpdateTest();
- }
+  void doGlobalUpdateTest() throws Exception {
+    runCommandForced("--type=recurring_billing", "--timestamp=1984-12-18T00:00:00Z");
+    assertThat(
+            ofy()
+                .load()
+                .key(Cursor.createGlobalKey(CursorType.RECURRING_BILLING))
+                .now()
+                .getCursorTime())
+        .isEqualTo(DateTime.parse("1984-12-18TZ"));
+  }
 
   @Test
-  public void testUpdateCursors_hasOldValue() throws Exception {
+  public void testSuccess_oldValueIsAbsent() throws Exception {
+    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry)).now()).isNull();
+    doUpdateTest();
+  }
+
+  @Test
+  public void testSuccess_hasOldValue() throws Exception {
     persistResource(Cursor.create(CursorType.BRDA, DateTime.parse("1950-12-18TZ"), registry));
     doUpdateTest();
   }
+
+  @Test
+  public void testSuccess_global_hasOldValue() throws Exception {
+    persistResource(
+        Cursor.createGlobal(CursorType.RECURRING_BILLING, DateTime.parse("1950-12-18TZ")));
+    doGlobalUpdateTest();
+  }
+
+  @Test
+  public void testSuccess_global_oldValueIsAbsent() throws Exception {
+    assertThat(ofy().load().key(Cursor.createGlobalKey(CursorType.RECURRING_BILLING)).now())
+        .isNull();
+    doGlobalUpdateTest();
+  }
+
+  @Test
+  public void testSuccess_multipleTlds_hasOldValue() throws Exception {
+    createTld("bar");
+    Registry registry2 = Registry.get("bar");
+    persistResource(Cursor.create(CursorType.BRDA, DateTime.parse("1950-12-18TZ"), registry));
+    persistResource(Cursor.create(CursorType.BRDA, DateTime.parse("1950-12-18TZ"), registry2));
+    runCommandForced("--type=brda", "--timestamp=1984-12-18T00:00:00Z", "foo", "bar");
+    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry)).now().getCursorTime())
+        .isEqualTo(DateTime.parse("1984-12-18TZ"));
+    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry2)).now().getCursorTime())
+        .isEqualTo(DateTime.parse("1984-12-18TZ"));
+  }
+
+  @Test
+  public void testSuccess_multipleTlds_oldValueIsAbsent() throws Exception {
+    createTld("bar");
+    Registry registry2 = Registry.get("bar");
+    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry)).now()).isNull();
+    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry2)).now()).isNull();
+    runCommandForced("--type=brda", "--timestamp=1984-12-18T00:00:00Z", "foo", "bar");
+    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry)).now().getCursorTime())
+        .isEqualTo(DateTime.parse("1984-12-18TZ"));
+    assertThat(ofy().load().key(Cursor.createKey(CursorType.BRDA, registry2)).now().getCursorTime())
+        .isEqualTo(DateTime.parse("1984-12-18TZ"));
+  }
+
+  @Test
+  public void testFailure_badTld() throws Exception {
+    thrown.expect(RegistryNotFoundException.class);
+    runCommandForced("--type=brda", "--timestamp=1984-12-18T00:00:00Z", "bar");
+  }
+
+  @Test
+  public void testFailure_badCursorType() throws Exception {
+    thrown.expect(ParameterException.class, "Invalid value for --type parameter");
+    runCommandForced("--type=rbda", "--timestamp=1984-12-18T00:00:00Z", "foo");
+  }
 }