Add a "loadtest" command to nomulus tool

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

java/google/registry/loadtest/LoadTestAction.java

@@ -53,7 +53,7 @@ import org.joda.time.DateTime;
  * least one must be specified in order for load testing to do anything.
  */
 @Action(
-    path = "/_dr/loadtest",
+    path = LoadTestAction.PATH,
     method = Action.Method.POST,
     automaticallyPrintOk = true)
 public class LoadTestAction implements Runnable {
@@ -71,6 +71,8 @@ public class LoadTestAction implements Runnable {
 
   private static final Random random = new Random();
 
+  public static final String PATH = "/_dr/loadtest";
+
   /** The client identifier of the registrar to use for load testing. */
   @Inject
   @Parameter("loadtestClientId")

java/google/registry/tools/BUILD

@@ -44,6 +44,7 @@ java_library(
         "//java/google/registry/flows",
         "//java/google/registry/gcs",
         "//java/google/registry/keyring/api",
+        "//java/google/registry/loadtest",
         "//java/google/registry/model",
         "//java/google/registry/pricing",
         "//java/google/registry/rde",

java/google/registry/tools/LoadTestCommand.java

@@ -0,0 +1,135 @@
+// Copyright 2017 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 com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.net.MediaType;
+import google.registry.loadtest.LoadTestAction;
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registry.Registries;
+
+/** Command to initiate a load-test. */
+@Parameters(separators = " =", commandDescription = "Run a load test.")
+class LoadTestCommand extends ConfirmingCommand implements ServerSideCommand {
+
+  // This is a mostly arbitrary value, roughly an hour and a quarter.  It served as a generous
+  // timespan for initial backup/restore testing, but has no other special significance.
+  private static final int DEFAULT_RUN_SECONDS = 4600;
+
+  @Parameter(
+      names = {"--tld"},
+      description = "TLD that all records will be created under.")
+  String tld = "example";
+
+  @Parameter(
+      names = {"--client_id"},
+      description = "Client ID of the registrar that will own new records.")
+  // "acme" is the id of the registrar we create in our public setup examples.
+  String clientId = "acme";
+
+  @Parameter(
+      names = {"--successful_host_creates"},
+      description = "Number of hosts to create per second.")
+  int successfulHostCreates = 1;
+
+  @Parameter(
+      names = {"--successful_domain_creates"},
+      description = "Number of domains to create per second.")
+  int successfulDomainCreates = 1;
+
+  @Parameter(
+      names = {"--successful_contact_creates"},
+      description = "Number of contact records to create per second.")
+  int successfulContactCreates = 1;
+
+  @Parameter(
+      names = {"--host_infos"},
+      description = "Number of successful host:info commands to send per second.")
+  int hostInfos = 1;
+
+  @Parameter(
+      names = {"--domain_infos"},
+      description = "Number of successful domain:info commands to send per second.")
+  int domainInfos = 1;
+
+  @Parameter(
+      names = {"--contact_infos"},
+      description = "Number of successful contact:info commands to send per second.")
+  int contactInfos = 1;
+
+  @Parameter(
+      names = {"--run_seconds"},
+      description = "Time to run the load test in seconds.")
+  int runSeconds = DEFAULT_RUN_SECONDS;
+
+  private Connection connection;
+
+  @Override
+  public void setConnection(Connection connection) {
+    this.connection = connection;
+  }
+
+  @Override
+  protected boolean checkExecutionState() throws Exception {
+    if (RegistryToolEnvironment.get() == RegistryToolEnvironment.PRODUCTION) {
+      System.err.println("You may not run a load test against production.");
+      return false;
+    }
+
+    // Check validity of TLD and Client Id.
+    if (!Registries.getTlds().contains(tld)) {
+      System.err.println("No such TLD: " + tld);
+      return false;
+    }
+    if (Registrar.loadByClientId(clientId) == null) {
+      System.err.println("No such client: " + clientId);
+      return false;
+    }
+
+    return true;
+  }
+
+  @Override
+  protected String prompt() {
+    return String.format(
+        "Run the load test (TLD = %s, Registry = %s, env = %s)?",
+        tld, clientId, RegistryToolEnvironment.get());
+  }
+
+  @Override
+  protected String execute() throws Exception {
+    System.err.println("Initiating load test...");
+
+    ImmutableMap<String, Object> params = new ImmutableMap.Builder<String, Object>()
+        .put("tld", tld)
+        .put("clientId", clientId)
+        .put("successfulHostCreates", successfulHostCreates)
+        .put("successfulDomainCreates", successfulDomainCreates)
+        .put("successfulContactCreates", successfulContactCreates)
+        .put("hostInfos", hostInfos)
+        .put("domainInfos", domainInfos)
+        .put("contactInfos", contactInfos)
+        .put("runSeconds", runSeconds)
+        .build();
+
+    return connection.send(
+        LoadTestAction.PATH,
+        params,
+        MediaType.PLAIN_TEXT_UTF_8,
+        new byte[0]);
+  }
+}

java/google/registry/tools/RegistryTool.java

@@ -89,6 +89,7 @@ public final class RegistryTool {
           .put("list_reserved_lists", ListReservedListsCommand.class)
           .put("list_tlds", ListTldsCommand.class)
           .put("load_snapshot", LoadSnapshotCommand.class)
+          .put("load_test", LoadTestCommand.class)
           .put("login", LoginCommand.class)
           .put("logout", LogoutCommand.class)
           .put("make_billing_tables", MakeBillingTablesCommand.class)

javatests/google/registry/tools/CommandTestCase.java

@@ -182,6 +182,10 @@ public abstract class CommandTestCase<C extends Command> {
     return new String(stdout.toByteArray(), UTF_8);
   }
 
+  protected String getStderrAsString() {
+    return new String(stderr.toByteArray(), UTF_8);
+  }
+
   protected List<String> getStdoutAsLines() {
     return Splitter.on('\n').omitEmptyStrings().trimResults().splitToList(getStdoutAsString());
   }

javatests/google/registry/tools/LoadTestCommandTest.java

@@ -0,0 +1,124 @@
+// Copyright 2017 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.createTld;
+import static google.registry.testing.DatastoreHelper.persistNewRegistrar;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.net.MediaType;
+import google.registry.testing.ExceptionRule;
+import google.registry.tools.ServerSideCommand.Connection;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class LoadTestCommandTest extends CommandTestCase<LoadTestCommand> {
+  Connection connection = mock(Connection.class);
+
+  @Rule public final ExceptionRule thrown = new ExceptionRule();
+
+  @Before
+  public void setUp() throws Exception {
+    command.setConnection(connection);
+    createTld("example");
+    persistNewRegistrar("acme", 99);
+  }
+
+  @Test
+  public void test_defaults() throws Exception {
+    runCommandForced();
+    ImmutableMap<String, Object> parms = new ImmutableMap.Builder<String, Object>()
+        .put("tld", "example")
+        .put("clientId", "acme")
+        .put("successfulHostCreates", 1)
+        .put("successfulDomainCreates", 1)
+        .put("successfulContactCreates", 1)
+        .put("hostInfos", 1)
+        .put("domainInfos", 1)
+        .put("contactInfos", 1)
+        .put("runSeconds", 4600)
+        .build();
+    verify(connection).send(
+        eq("/_dr/loadtest"),
+        eq(parms),
+        eq(MediaType.PLAIN_TEXT_UTF_8),
+        eq(new byte[0]));
+  }
+
+  @Test
+  public void test_overrides() throws Exception {
+    createTld("foo");
+    runCommandForced(
+        "--tld=foo",
+        "--client_id=NewRegistrar",
+        "--successful_host_creates=10",
+        "--successful_domain_creates=11",
+        "--successful_contact_creates=12",
+        "--host_infos=13",
+        "--domain_infos=14",
+        "--contact_infos=15",
+        "--run_seconds=16");
+    ImmutableMap<String, Object> parms = new ImmutableMap.Builder<String, Object>()
+        .put("tld", "foo")
+        .put("clientId", "NewRegistrar")
+        .put("successfulHostCreates", 10)
+        .put("successfulDomainCreates", 11)
+        .put("successfulContactCreates", 12)
+        .put("hostInfos", 13)
+        .put("domainInfos", 14)
+        .put("contactInfos", 15)
+        .put("runSeconds", 16)
+        .build();
+    verify(connection).send(
+        eq("/_dr/loadtest"),
+        eq(parms),
+        eq(MediaType.PLAIN_TEXT_UTF_8),
+        eq(new byte[0]));
+  }
+  @Test
+  public void test_prompt() throws Exception {
+    thrown.expect(IllegalStateException.class);
+    runCommand();
+    verifyZeroInteractions(connection);
+    assertInStderr("Unable to access stdin (are you running with bazel run?)");
+  }
+
+  @Test
+  public void test_badTLD() throws Exception {
+    runCommand("--tld=bogus");
+    verifyZeroInteractions(connection);
+    assertInStderr("No such TLD: bogus");
+  }
+
+  @Test
+  public void test_badClientId() throws Exception {
+    runCommand("--client_id=badaddry");
+    verifyZeroInteractions(connection);
+    assertInStderr("No such client: badaddry");
+  }
+
+  @Test
+  public void test_noProduction() throws Exception {
+    runCommandInEnvironment(RegistryToolEnvironment.PRODUCTION);
+    assertInStderr("You may not run a load test against production.");
+  }
+}