From 91febfd94588317ca8b40241b4e94b6811c7d0c7 Mon Sep 17 00:00:00 2001 From: mcilwain Date: Thu, 21 Mar 2019 10:36:58 -0700 Subject: [PATCH] Add a nomulus tool command to get allocation tokens ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=239624232 --- .../tools/GetAllocationTokenCommand.java | 57 ++++++++++++++++ java/google/registry/tools/GtechTool.java | 1 + java/google/registry/tools/RegistryTool.java | 1 + .../tools/GetAllocationTokenCommandTest.java | 65 +++++++++++++++++++ 4 files changed, 124 insertions(+) create mode 100644 java/google/registry/tools/GetAllocationTokenCommand.java create mode 100644 javatests/google/registry/tools/GetAllocationTokenCommandTest.java diff --git a/java/google/registry/tools/GetAllocationTokenCommand.java b/java/google/registry/tools/GetAllocationTokenCommand.java new file mode 100644 index 000000000..ab1e507cc --- /dev/null +++ b/java/google/registry/tools/GetAllocationTokenCommand.java @@ -0,0 +1,57 @@ +// Copyright 2019 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static google.registry.model.ofy.ObjectifyService.ofy; + +import com.beust.jcommander.Parameter; +import com.beust.jcommander.Parameters; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Lists; +import com.googlecode.objectify.Key; +import google.registry.model.domain.token.AllocationToken; +import java.util.List; + +/** Command to show allocation tokens. */ +@Parameters(separators = " =", commandDescription = "Show allocation token(s)") +final class GetAllocationTokenCommand implements CommandWithRemoteApi { + + @Parameter( + description = "Allocation token(s)", + required = true) + private List mainParameters; + + private static final int BATCH_SIZE = 20; + + @Override + public void run() { + ImmutableMap.Builder builder = new ImmutableMap.Builder<>(); + for (List tokens : Lists.partition(mainParameters, BATCH_SIZE)) { + ImmutableList> keys = + tokens.stream().map(t -> Key.create(AllocationToken.class, t)).collect(toImmutableList()); + ofy().load().keys(keys).forEach((k, v) -> builder.put(k.getName(), v)); + } + ImmutableMap loadedTokens = builder.build(); + + for (String token : mainParameters) { + System.out.println( + loadedTokens.containsKey(token) + ? loadedTokens.get(token) + : String.format("Token %s does not exist.", token)); + } + } +} diff --git a/java/google/registry/tools/GtechTool.java b/java/google/registry/tools/GtechTool.java index 90f0cce9a..427c29bb9 100644 --- a/java/google/registry/tools/GtechTool.java +++ b/java/google/registry/tools/GtechTool.java @@ -40,6 +40,7 @@ public final class GtechTool { "create_sandbox_tld", "delete_domain", "generate_dns_report", + "get_allocation_token", "get_claims_list", "get_contact", "get_domain", diff --git a/java/google/registry/tools/RegistryTool.java b/java/google/registry/tools/RegistryTool.java index 460d159f6..5beb36bd3 100644 --- a/java/google/registry/tools/RegistryTool.java +++ b/java/google/registry/tools/RegistryTool.java @@ -61,6 +61,7 @@ public final class RegistryTool { .put("generate_escrow_deposit", GenerateEscrowDepositCommand.class) .put("generate_lordn", GenerateLordnCommand.class) .put("generate_zone_files", GenerateZoneFilesCommand.class) + .put("get_allocation_token", GetAllocationTokenCommand.class) .put("get_claims_list", GetClaimsListCommand.class) .put("get_contact", GetContactCommand.class) .put("get_domain", GetDomainCommand.class) diff --git a/javatests/google/registry/tools/GetAllocationTokenCommandTest.java b/javatests/google/registry/tools/GetAllocationTokenCommandTest.java new file mode 100644 index 000000000..8b8fd0906 --- /dev/null +++ b/javatests/google/registry/tools/GetAllocationTokenCommandTest.java @@ -0,0 +1,65 @@ +// Copyright 2019 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static google.registry.testing.DatastoreHelper.persistResource; +import static google.registry.testing.DatastoreHelper.persistSimpleResources; +import static google.registry.testing.JUnitBackports.assertThrows; + +import com.beust.jcommander.ParameterException; +import com.google.common.collect.ImmutableList; +import google.registry.model.domain.token.AllocationToken; +import org.joda.time.DateTime; +import org.junit.Test; + +/** Unit tests for {@link GetAllocationTokenCommand}. */ +public class GetAllocationTokenCommandTest extends CommandTestCase { + + @Test + public void testSuccess_oneToken() throws Exception { + AllocationToken token = + persistResource( + new AllocationToken.Builder().setToken("foo").setDomainName("foo.bar").build()); + runCommand("foo"); + assertInStdout(token.toHydratedString()); + } + + @Test + public void testSuccess_multipleTokens() throws Exception { + ImmutableList tokens = persistSimpleResources( + ImmutableList.of( + new AllocationToken.Builder() + .setToken("fee") + .setCreationTimeForTest(DateTime.parse("2015-04-07T22:19:17.044Z")) + .build(), + new AllocationToken.Builder().setToken("fii").setDomainName("bar.baz").build())); + runCommand("fee", "fii"); + assertInStdout(tokens.get(0).toHydratedString(), tokens.get(1).toHydratedString()); + } + + @Test + public void testSuccess_oneTokenDoesNotExist() throws Exception { + AllocationToken token = + persistResource( + new AllocationToken.Builder().setToken("foo").setDomainName("foo.bar").build()); + runCommand("foo", "bar"); + assertInStdout(token.toHydratedString(), "Token bar does not exist."); + } + + @Test + public void testFailure_noAllocationTokensSpecified() { + assertThrows(ParameterException.class, this::runCommand); + } +}