diff --git a/java/google/registry/tools/BUILD b/java/google/registry/tools/BUILD index 58349dcbc..45fc9b432 100644 --- a/java/google/registry/tools/BUILD +++ b/java/google/registry/tools/BUILD @@ -66,6 +66,7 @@ java_library( "//java/google/registry/xjc", "//java/google/registry/xml", "//third_party/java/appengine:appengine-api", + "//third_party/java/appengine:appengine-remote-api-link", "//third_party/java/bouncycastle", "//third_party/java/bouncycastle_bcpg", "//third_party/java/dagger", diff --git a/java/google/registry/tools/CreateLrpTokensCommand.java b/java/google/registry/tools/CreateLrpTokensCommand.java index 61802a31a..2d52242fb 100644 --- a/java/google/registry/tools/CreateLrpTokensCommand.java +++ b/java/google/registry/tools/CreateLrpTokensCommand.java @@ -24,6 +24,8 @@ import static java.nio.charset.StandardCharsets.UTF_8; import com.beust.jcommander.Parameter; import com.beust.jcommander.Parameters; +import com.google.appengine.tools.remoteapi.RemoteApiException; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.CharMatcher; import com.google.common.base.Function; import com.google.common.base.Splitter; @@ -40,24 +42,28 @@ import google.registry.tools.Command.RemoteApiCommand; import google.registry.tools.params.KeyValueMapParameter.StringToIntegerMap; import google.registry.tools.params.KeyValueMapParameter.StringToStringMap; import google.registry.tools.params.PathParameter; +import google.registry.util.NonFinalForTesting; +import google.registry.util.Retrier; import google.registry.util.StringGenerator; import google.registry.util.TokenUtils; import java.io.StringReader; import java.nio.file.Path; import java.util.Collection; import java.util.Set; +import java.util.concurrent.Callable; import javax.inject.Inject; /** * Command to create one or more LRP tokens, given assignee(s) as either a parameter or a text file. */ +@NonFinalForTesting @Parameters( separators = " =", commandDescription = "Create an LRP token for a given assignee (using -a) or import a text" + " file of assignees for bulk token creation (using -i). Assignee/token pairs are printed" + " to stdout, and should be piped to a file for distribution to assignees or for cleanup" + " in the event of a command interruption.") -public final class CreateLrpTokensCommand implements RemoteApiCommand { +public class CreateLrpTokensCommand implements RemoteApiCommand { @Parameter( names = {"-a", "--assignee"}, @@ -95,6 +101,7 @@ public final class CreateLrpTokensCommand implements RemoteApiCommand { private ImmutableMap metadataColumns; @Inject StringGenerator stringGenerator; + @Inject Retrier retrier; private static final int BATCH_SIZE = 20; @@ -126,7 +133,7 @@ public final class CreateLrpTokensCommand implements RemoteApiCommand { String line = null; do { - ImmutableSet.Builder tokensToSave = new ImmutableSet.Builder<>(); + ImmutableSet.Builder tokensToSaveBuilder = new ImmutableSet.Builder<>(); for (String token : generateTokens(BATCH_SIZE)) { line = reader.readLine(); if (!isNullOrEmpty(line)) { @@ -155,14 +162,22 @@ public final class CreateLrpTokensCommand implements RemoteApiCommand { } tokenBuilder.setMetadata(metadataBuilder.build()); } - tokensToSave.add(tokenBuilder.build()); + tokensToSaveBuilder.add(tokenBuilder.build()); } } - saveTokens(tokensToSave.build()); + final ImmutableSet tokensToSave = tokensToSaveBuilder.build(); + // Wrap in a retrier to deal with transient 404 errors (thrown as RemoteApiExceptions). + retrier.callWithRetry(new Callable() { + @Override + public Void call() throws Exception { + saveTokens(tokensToSave); + return null; + }}, RemoteApiException.class); } while (line != null); } - private void saveTokens(final ImmutableSet tokens) { + @VisibleForTesting + void saveTokens(final ImmutableSet tokens) { Collection savedTokens = ofy().transact(new Work>() { @Override diff --git a/java/google/registry/tools/RegistryToolComponent.java b/java/google/registry/tools/RegistryToolComponent.java index c92262de4..13d62a4da 100644 --- a/java/google/registry/tools/RegistryToolComponent.java +++ b/java/google/registry/tools/RegistryToolComponent.java @@ -25,6 +25,7 @@ import google.registry.request.Modules.DatastoreServiceModule; import google.registry.request.Modules.Jackson2Module; import google.registry.request.Modules.URLFetchServiceModule; import google.registry.util.SystemClock.SystemClockModule; +import google.registry.util.SystemSleeper.SystemSleeperModule; /** * Dagger component for Registry Tool. @@ -43,6 +44,7 @@ import google.registry.util.SystemClock.SystemClockModule; KeyModule.class, RegistryToolModule.class, SystemClockModule.class, + SystemSleeperModule.class, URLFetchServiceModule.class, VoidDnsWriterModule.class, } diff --git a/java/google/registry/util/NonFinalForTesting.java b/java/google/registry/util/NonFinalForTesting.java index aeb3abdf3..848c65fc2 100644 --- a/java/google/registry/util/NonFinalForTesting.java +++ b/java/google/registry/util/NonFinalForTesting.java @@ -16,6 +16,7 @@ package google.registry.util; import static java.lang.annotation.ElementType.FIELD; import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.TYPE; import static java.lang.annotation.RetentionPolicy.SOURCE; import java.lang.annotation.Documented; @@ -36,5 +37,5 @@ import java.lang.annotation.Target; */ @Documented @Retention(SOURCE) -@Target({FIELD, METHOD}) +@Target({FIELD, METHOD, TYPE}) public @interface NonFinalForTesting {} diff --git a/javatests/google/registry/tools/BUILD b/javatests/google/registry/tools/BUILD index 6117e6887..8b8a11868 100644 --- a/javatests/google/registry/tools/BUILD +++ b/javatests/google/registry/tools/BUILD @@ -24,6 +24,7 @@ java_library( "//java/com/google/common/net", "//java/com/google/common/reflect", "//third_party/java/appengine:appengine-api-testonly", + "//third_party/java/appengine:appengine-remote-api-link", "//third_party/java/jcommander", "//third_party/java/joda_money", "//third_party/java/joda_time", diff --git a/javatests/google/registry/tools/CreateLrpTokensCommandTest.java b/javatests/google/registry/tools/CreateLrpTokensCommandTest.java index 84c2af85b..cd1447040 100644 --- a/javatests/google/registry/tools/CreateLrpTokensCommandTest.java +++ b/javatests/google/registry/tools/CreateLrpTokensCommandTest.java @@ -19,7 +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 static java.nio.charset.StandardCharsets.UTF_8; +import static org.mockito.Matchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.spy; +import com.google.appengine.tools.remoteapi.RemoteApiException; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.io.Files; @@ -28,10 +32,14 @@ import google.registry.model.domain.LrpTokenEntity; import google.registry.model.reporting.HistoryEntry; import google.registry.testing.DeterministicStringGenerator; import google.registry.testing.DeterministicStringGenerator.Rule; +import google.registry.testing.FakeClock; +import google.registry.testing.FakeSleeper; +import google.registry.util.Retrier; import java.io.File; import java.io.IOException; import java.util.Set; import javax.annotation.Nullable; +import org.joda.time.DateTime; import org.junit.Before; import org.junit.Test; @@ -48,6 +56,8 @@ public class CreateLrpTokensCommandTest extends CommandTestCase