Support datastore restore in Nomulus tool

Two commands are being added:
- ImportDatastoreCommand starts an async import operation.
  User may choose to wait until import completes or quit
  immediately.
- GetOperationStatusCommand checks the status of an operation.
  It may be used to check the status of an operation started by
  ImportDatastoreCommand.

Both commands communicate with Datastore admin api directly, without
going through the Registry server.

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

java/google/registry/export/CheckBackupAction.java

@@ -14,6 +14,7 @@
 
 package google.registry.export;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Sets.intersection;
 import static google.registry.export.UploadDatastoreBackupAction.enqueueUploadBackupTask;
 import static google.registry.request.Action.Method.GET;
@@ -131,6 +132,8 @@ public class CheckBackupAction implements Runnable {
     Set<String> kindsToLoad = ImmutableSet.copyOf(Splitter.on(',').split(kindsToLoadParam));
     Operation backup = getExportStatus();
 
+    checkArgument(backup.isExport(), "Expecting an export operation: [%s].", backupName);
+
     if (backup.isProcessing()
         && backup.getRunningTime(clock).isShorterThan(MAXIMUM_BACKUP_RUNNING_TIME)) {
       // Backup might still be running, so send a 304 to have the task retry.

java/google/registry/export/datastore/DatastoreAdmin.java

@@ -25,6 +25,7 @@ import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonFactory;
 import com.google.api.client.util.Key;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
 import java.util.Collection;
 import java.util.Optional;
 
@@ -73,12 +74,32 @@ public class DatastoreAdmin extends AbstractGoogleJsonClient {
    * </ul>
    *
    * @param outputUrlPrefix the full resource URL of the external storage location
-   * @param kinds the datastore 'kinds' to be exported
+   * @param kinds the datastore 'kinds' to be exported. If empty, all kinds will be exported
    */
   public Export export(String outputUrlPrefix, Collection<String> kinds) {
     return new Export(new ExportRequest(outputUrlPrefix, kinds));
   }
 
+  /**
+   * Imports the entire backup specified by {@code backupUrl} back to Cloud Datastore.
+   *
+   * <p>A successful backup restores deleted entities and reverts updates to existing entities since
+   * the backup time. However, it does not affect newly added entities.
+   */
+  public Import importBackup(String backupUrl) {
+    return new Import(new ImportRequest(backupUrl, ImmutableList.of()));
+  }
+
+  /**
+   * Imports the backup specified by {@code backupUrl} back to Cloud Datastore. Only entities whose
+   * types are included in {@code kinds} are imported.
+   *
+   * @see #importBackup(String)
+   */
+  public Import importBackup(String backupUrl, Collection<String> kinds) {
+    return new Import(new ImportRequest(backupUrl, kinds));
+  }
+
   /**
    * Returns a {@link Get} request that retrieves the details of an export or import {@link
    * Operation}.
@@ -158,6 +179,20 @@ public class DatastoreAdmin extends AbstractGoogleJsonClient {
     }
   }
 
+  /** A request to restore an backup to a Cloud Datastore database. */
+  public class Import extends DatastoreAdminRequest<Operation> {
+
+    Import(ImportRequest importRequest) {
+      super(
+          DatastoreAdmin.this,
+          "POST",
+          "projects/{projectId}:import",
+          importRequest,
+          Operation.class);
+      set("projectId", projectId);
+    }
+  }
+
   /** A request to retrieve details of an export or import operation. */
   public class Get extends DatastoreAdminRequest<Operation> {
 
@@ -216,8 +251,27 @@ public class DatastoreAdmin extends AbstractGoogleJsonClient {
 
     ExportRequest(String outputUrlPrefix, Collection<String> kinds) {
       checkNotNull(outputUrlPrefix, "outputUrlPrefix");
+      checkArgument(!kinds.isEmpty(), "kinds must not be empty");
       this.outputUrlPrefix = outputUrlPrefix;
       this.entityFilter = new EntityFilter(kinds);
     }
   }
+
+  /**
+   * Model object that describes the JSON content in an export request.
+   *
+   * <p>Please note that some properties defined in the API are excluded, e.g., {@code databaseId}
+   * (not supported by Cloud Datastore) and labels (not used by Domain Registry).
+   */
+  @SuppressWarnings("unused")
+  static class ImportRequest extends GenericJson {
+    @Key private final String inputUrl;
+    @Key private final EntityFilter entityFilter;
+
+    ImportRequest(String inputUrl, Collection<String> kinds) {
+      checkNotNull(inputUrl, "outputUrlPrefix");
+      this.inputUrl = inputUrl;
+      this.entityFilter = new EntityFilter(kinds);
+    }
+  }
 }

java/google/registry/export/datastore/EntityFilter.java

@@ -14,7 +14,6 @@
 
 package google.registry.export.datastore;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.api.client.json.GenericJson;
@@ -39,7 +38,6 @@ public class EntityFilter extends GenericJson {
 
   EntityFilter(Collection<String> kinds) {
     checkNotNull(kinds, "kinds");
-    checkArgument(!kinds.isEmpty(), "kinds must not be empty");
     this.kinds = ImmutableList.copyOf(kinds);
   }
 

java/google/registry/export/datastore/Operation.java

@@ -15,10 +15,10 @@
 package google.registry.export.datastore;
 
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.base.Strings.isNullOrEmpty;
 
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.util.Key;
-import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import google.registry.export.datastore.DatastoreAdmin.Get;
@@ -52,6 +52,14 @@ public class Operation extends GenericJson {
     return name;
   }
 
+  public boolean isExport() {
+    return !isNullOrEmpty(getExportFolderUrl());
+  }
+
+  public boolean isImport() {
+    return !isNullOrEmpty(getMetadata().getInputUrl());
+  }
+
   public boolean isDone() {
     return done;
   }
@@ -88,7 +96,9 @@ public class Operation extends GenericJson {
   /**
    * Returns the URL to the GCS folder that holds the exported data. This folder is created by
    * Datastore and is under the {@code outputUrlPrefix} set to {@linkplain
-   * DatastoreAdmin#export(String, List) the export request}.
+   * DatastoreAdmin#export(String, java.util.Collection) the export request}.
+   *
+   * @throws IllegalStateException if this is not an export operation
    */
   public String getExportFolderUrl() {
     return getMetadata().getOutputUrlPrefix();
@@ -98,6 +108,8 @@ public class Operation extends GenericJson {
    * Returns the last segment of the {@linkplain #getExportFolderUrl() export folder URL} which can
    * be used as unique identifier of this export operation. This is a better ID than the {@linkplain
    * #getName() operation name}, which is opaque.
+   *
+   * @throws IllegalStateException if this is not an export operation
    */
   public String getExportId() {
     String exportFolderUrl = getExportFolderUrl();
@@ -138,12 +150,12 @@ public class Operation extends GenericJson {
     public CommonMetadata() {}
 
     String getOperationType() {
-      checkState(!Strings.isNullOrEmpty(operationType), "operationType may not be null or empty");
+      checkState(!isNullOrEmpty(operationType), "operationType may not be null or empty");
       return operationType;
     }
 
     String getState() {
-      checkState(!Strings.isNullOrEmpty(state), "state may not be null or empty");
+      checkState(!isNullOrEmpty(state), "state may not be null or empty");
       return state;
     }
 
@@ -165,6 +177,7 @@ public class Operation extends GenericJson {
     @Key private Progress progressEntities;
     @Key private Progress progressBytes;
     @Key private EntityFilter entityFilter;
+    @Key private String inputUrl;
     @Key private String outputUrlPrefix;
 
     public Metadata() {}
@@ -186,9 +199,22 @@ public class Operation extends GenericJson {
       return entityFilter;
     }
 
+    public String getInputUrl() {
+      return checkUrls().inputUrl;
+    }
+
     public String getOutputUrlPrefix() {
-      checkState(!Strings.isNullOrEmpty(outputUrlPrefix), "outputUrlPrefix");
-      return outputUrlPrefix;
+      return checkUrls().outputUrlPrefix;
+    }
+
+    Metadata checkUrls() {
+      checkState(
+          isNullOrEmpty(inputUrl) || isNullOrEmpty(outputUrlPrefix),
+          "inputUrl and outputUrlPrefix must not be both present");
+      checkState(
+          !isNullOrEmpty(inputUrl) || !isNullOrEmpty(outputUrlPrefix),
+          "inputUrl and outputUrlPrefix must not be both missing");
+      return this;
     }
   }
 

java/google/registry/tools/BUILD

@@ -45,6 +45,7 @@ java_library(
         "//java/google/registry/dns/writer/clouddns",
         "//java/google/registry/dns/writer/dnsupdate",
         "//java/google/registry/export",
+        "//java/google/registry/export/datastore",
         "//java/google/registry/flows",
         "//java/google/registry/gcs",
         "//java/google/registry/keyring",

java/google/registry/tools/GetOperationStatusCommand.java

@@ -0,0 +1,47 @@
+// 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.base.Preconditions.checkArgument;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.api.client.json.JsonFactory;
+import com.google.api.client.json.jackson2.JacksonFactory;
+import com.google.common.base.Strings;
+import google.registry.export.datastore.DatastoreAdmin;
+import java.util.List;
+import javax.inject.Inject;
+
+/** Command to get the status of a Datastore operation, e.g., an import or export. */
+@Parameters(separators = " =", commandDescription = "Get status of a Datastore operation.")
+public class GetOperationStatusCommand implements Command {
+
+  private static final JsonFactory JSON_FACTORY = JacksonFactory.getDefaultInstance();
+
+  @Parameter(description = "Name of the Datastore import or export operation.")
+  private List<String> mainParameters;
+
+  @Inject DatastoreAdmin datastoreAdmin;
+
+  @Override
+  public void run() throws Exception {
+    checkArgument(
+        mainParameters.size() == 1, "Requires exactly one argument: the name of the operation.");
+    String operationName = mainParameters.get(0);
+    checkArgument(!Strings.isNullOrEmpty(operationName), "Missing operation name.");
+    System.out.println(JSON_FACTORY.toPrettyString(datastoreAdmin.get(operationName).execute()));
+  }
+}

java/google/registry/tools/ImportDatastoreCommand.java

@@ -0,0 +1,127 @@
+// 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.base.Preconditions.checkArgument;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.base.Ascii;
+import com.google.common.collect.ImmutableList;
+import google.registry.export.datastore.DatastoreAdmin;
+import google.registry.export.datastore.Operation;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import javax.inject.Inject;
+import org.joda.time.Duration;
+
+/**
+ * Command that imports an earlier backup into Datastore.
+ *
+ * <p>This command is part of the Datastore restore process. Please refer to <a
+ * href="http://playbooks/domain_registry/procedures/backup-restore-testing.md">the playbook</a> for
+ * the entire process.
+ */
+@Parameters(separators = " =", commandDescription = "Imports a backup of the Datastore.")
+public class ImportDatastoreCommand extends ConfirmingCommand {
+
+  @Parameter(names = "--backup_url", description = "URL to the backup on GCS to be imported.")
+  private String backupUrl;
+
+  @Nullable
+  @Parameter(
+      names = "--kinds",
+      description = "List of entity kinds to be imported. Default is to import all.")
+  private List<String> kinds = ImmutableList.of();
+
+  @Parameter(
+      names = "--async",
+      description = "If true, command will launch import operation and quit.")
+  private boolean async;
+
+  @Parameter(
+      names = "--poll_interval",
+      description =
+          "Polling interval while waiting for completion synchronously. "
+              + "Value is in ISO-8601 format, e.g., PT10S for 10 seconds.")
+  private Duration pollingInterval = Duration.standardSeconds(30);
+
+  @Parameter(
+      names = "--confirm_production_import",
+      description = "Set this option to 'PRODUCTION' to confirm import in production environment.")
+  private String confirmProductionImport = "";
+
+  @Inject DatastoreAdmin datastoreAdmin;
+
+  @Override
+  protected String execute() throws Exception {
+    RegistryToolEnvironment currentEnvironment = RegistryToolEnvironment.get();
+
+    // Extra confirmation for running in production
+    checkArgument(
+        !currentEnvironment.equals(RegistryToolEnvironment.PRODUCTION)
+            || confirmProductionImport.equals("PRODUCTION"),
+        "The confirm_production_import option must be set when restoring production environment.");
+
+    Operation importOperation = datastoreAdmin.importBackup(backupUrl, kinds).execute();
+
+    String statusCommand =
+        String.format(
+            "nomulus -e %s get_operation_status %s",
+            Ascii.toLowerCase(currentEnvironment.name()), importOperation.getName());
+
+    if (async) {
+      return String.format(
+          "Datastore import started. Run this command to check its progress:\n%s",
+          statusCommand);
+    }
+
+    System.out.println(
+        "Waiting for import to complete.\n"
+            + "You may press Ctrl-C at any time, and use this command to check progress:\n"
+            + statusCommand);
+    while (importOperation.isProcessing()) {
+      waitInteractively(pollingInterval);
+
+      importOperation = datastoreAdmin.get(importOperation.getName()).execute();
+
+      System.out.printf("\n%s\n", importOperation.getProgress());
+    }
+    return String.format(
+        "\nDatastore import %s %s.",
+        importOperation.getName(), importOperation.isSuccessful() ? "succeeded" : "failed");
+  }
+
+  @Override
+  protected String prompt() {
+    return "\nThis command is an intermediate step in the Datastore restore process.\n\n"
+        + "Please read and understand the playbook entry at\n"
+        + "    http://playbooks/domain_registry/procedures/backup-restore-testing.md\n"
+        + "before proceeding.\n";
+  }
+
+  /** Prints dots to console at regular interval while waiting. */
+  private static void waitInteractively(Duration pollingInterval) throws InterruptedException {
+    int sleepSeconds = 2;
+    long iterations = (pollingInterval.getStandardSeconds() + sleepSeconds - 1) / sleepSeconds;
+
+    for (int i = 0; i < iterations; i++) {
+      TimeUnit.SECONDS.sleep(sleepSeconds);
+      System.out.print('.');
+      System.out.flush();
+    }
+  }
+}

java/google/registry/tools/RegistryTool.java

@@ -75,6 +75,7 @@ public final class RegistryTool {
           .put("get_history_entries", GetHistoryEntriesCommand.class)
           .put("get_host", GetHostCommand.class)
           .put("get_keyring_secret", GetKeyringSecretCommand.class)
+          .put("get_operation_status", GetOperationStatusCommand.class)
           .put("get_registrar", GetRegistrarCommand.class)
           .put("get_resource_by_key", GetResourceByKeyCommand.class)
           .put("get_routing_map", GetRoutingMapCommand.class)
@@ -83,6 +84,7 @@ public final class RegistryTool {
           .put("get_tld", GetTldCommand.class)
           .put("ghostryde", GhostrydeCommand.class)
           .put("hash_certificate", HashCertificateCommand.class)
+          .put("import_datastore", ImportDatastoreCommand.class)
           .put("list_cursors", ListCursorsCommand.class)
           .put("list_domains", ListDomainsCommand.class)
           .put("list_hosts", ListHostsCommand.class)

java/google/registry/tools/RegistryToolComponent.java

@@ -22,6 +22,7 @@ import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.dns.writer.VoidDnsWriterModule;
 import google.registry.dns.writer.clouddns.CloudDnsWriterModule;
 import google.registry.dns.writer.dnsupdate.DnsUpdateWriterModule;
+import google.registry.export.datastore.DatastoreAdminModule;
 import google.registry.keyring.KeyringModule;
 import google.registry.keyring.api.DummyKeyringModule;
 import google.registry.keyring.api.KeyModule;
@@ -56,6 +57,7 @@ import javax.inject.Singleton;
       BigqueryModule.class,
       ConfigModule.class,
       CloudDnsWriterModule.class,
+      DatastoreAdminModule.class,
       DatastoreServiceModule.class,
       DummyKeyringModule.class,
       DnsUpdateWriterModule.class,
@@ -92,7 +94,9 @@ interface RegistryToolComponent {
   void inject(GenerateDnsReportCommand command);
   void inject(GenerateEscrowDepositCommand command);
   void inject(GetKeyringSecretCommand command);
+  void inject(GetOperationStatusCommand command);
   void inject(GhostrydeCommand command);
+  void inject(ImportDatastoreCommand command);
   void inject(ListCursorsCommand command);
   void inject(LoadSnapshotCommand command);
   void inject(LockDomainCommand command);

javatests/google/registry/export/datastore/EntityFilterTest.java

@@ -37,11 +37,6 @@ public class EntityFilterTest {
     assertThrows(NullPointerException.class, () -> new EntityFilter(null));
   }
 
-  @Test
-  public void testEntityFilter_create_emptyKinds() {
-    assertThrows(IllegalArgumentException.class, () -> new EntityFilter(ImmutableList.of()));
-  }
-
   @Test
   public void testEntityFilter_marshall() throws IOException {
     EntityFilter entityFilter =

javatests/google/registry/tools/BUILD

@@ -19,6 +19,7 @@ java_library(
     deps = [
         "//java/google/registry/backup",
         "//java/google/registry/config",
+        "//java/google/registry/export/datastore",
         "//java/google/registry/flows",
         "//java/google/registry/keyring/api",
         "//java/google/registry/model",

javatests/google/registry/tools/GetOperationStatusCommandTest.java

@@ -0,0 +1,59 @@
+// 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.truth.Truth.assertThat;
+import static google.registry.testing.JUnitBackports.assertThrows;
+import static org.mockito.Mockito.when;
+
+import google.registry.export.datastore.DatastoreAdmin;
+import google.registry.export.datastore.DatastoreAdmin.Get;
+import google.registry.export.datastore.Operation;
+import java.io.IOException;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+
+/** Unit tests for {@link GetOperationStatusCommand}. */
+@RunWith(JUnit4.class)
+public class GetOperationStatusCommandTest extends CommandTestCase<GetOperationStatusCommand> {
+
+  @Mock private DatastoreAdmin datastoreAdmin;
+  @Mock private Get getRequest;
+  @Captor ArgumentCaptor<String> operationName;
+
+  @Before
+  public void setup() throws IOException {
+    command.datastoreAdmin = datastoreAdmin;
+
+    when(datastoreAdmin.get(operationName.capture())).thenReturn(getRequest);
+    when(getRequest.execute()).thenReturn(new Operation());
+  }
+
+  @Test
+  public void test_success() throws Exception {
+    runCommand("projects/project-id/operations/HASH");
+    assertThat(operationName.getValue()).isEqualTo("projects/project-id/operations/HASH");
+  }
+
+  @Test
+  public void test_failure_tooManyNames() {
+    assertThrows(IllegalArgumentException.class, () -> runCommand("a", "b"));
+  }
+}

javatests/google/registry/tools/ImportDatastoreCommandTest.java

@@ -0,0 +1,146 @@
+// 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.truth.Truth.assertThat;
+import static google.registry.testing.JUnitBackports.assertThrows;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import google.registry.export.datastore.DatastoreAdmin;
+import google.registry.export.datastore.DatastoreAdmin.Get;
+import google.registry.export.datastore.DatastoreAdmin.Import;
+import google.registry.export.datastore.Operation;
+import java.util.Collection;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+
+/** Unit tests for {@link ImportDatastoreCommand}. */
+@RunWith(JUnit4.class)
+public class ImportDatastoreCommandTest extends CommandTestCase<ImportDatastoreCommand> {
+
+  @Captor ArgumentCaptor<String> backupUrl;
+  @Captor ArgumentCaptor<Collection<String>> kinds;
+  @Captor ArgumentCaptor<String> operationName;
+
+  @Mock private DatastoreAdmin datastoreAdmin;
+  @Mock private Import importRequest;
+  @Mock private Get getRequest;
+  @Mock private Operation importOperation;
+  @Mock private Operation getOperation;
+
+  @Before
+  public void setup() throws Exception {
+    command.datastoreAdmin = datastoreAdmin;
+
+    when(datastoreAdmin.importBackup(backupUrl.capture(), kinds.capture()))
+        .thenReturn(importRequest);
+    when(importRequest.execute()).thenReturn(importOperation);
+    when(importOperation.getName()).thenReturn("opName");
+
+    when(datastoreAdmin.get(operationName.capture())).thenReturn(getRequest);
+    when(getRequest.execute()).thenReturn(getOperation);
+  }
+
+  @Test
+  public void test_importAllKinds_immediateSuccess() throws Exception {
+    runCommandForced(
+        "--poll_interval", "PT0.001S",
+        "--backup_url", "gs://bucket/export-id/export-id.overall_export_metadata");
+    assertThat(backupUrl.getValue())
+        .isEqualTo("gs://bucket/export-id/export-id.overall_export_metadata");
+    assertThat(kinds.getValue()).isEmpty();
+    verify(datastoreAdmin, never()).get(anyString());
+  }
+
+  @Test
+  public void test_importSomeKinds_immediateSuccess() throws Exception {
+    runCommandForced(
+        "--poll_interval",
+        "PT0.001S",
+        "--backup_url",
+        "gs://bucket/export-id/export-id.overall_export_metadata",
+        "--kinds",
+        "Registrar",
+        "--kinds",
+        "Registry");
+    assertThat(backupUrl.getValue())
+        .isEqualTo("gs://bucket/export-id/export-id.overall_export_metadata");
+    assertThat(kinds.getValue()).containsExactly("Registrar", "Registry");
+    verify(datastoreAdmin, never()).get(anyString());
+  }
+
+  @Test
+  public void test_delayedSuccess_sync() throws Exception {
+    when(importOperation.isProcessing()).thenReturn(true);
+    when(getOperation.isProcessing()).thenReturn(true).thenReturn(false);
+    when(getOperation.isSuccessful()).thenReturn(true);
+    runCommandForced(
+        "--poll_interval", "PT0.001S",
+        "--backup_url", "gs://bucket/export-id/export-id.overall_export_metadata");
+
+    verify(datastoreAdmin, times(1)).get("opName");
+  }
+
+  @Test
+  public void test_delayedSuccess_async() throws Exception {
+    when(importOperation.isProcessing()).thenReturn(false);
+    when(getOperation.isProcessing()).thenReturn(true).thenReturn(false);
+    when(getOperation.isSuccessful()).thenReturn(true);
+    runCommandForced(
+        "--poll_interval",
+        "PT0.001S",
+        "--backup_url",
+        "gs://bucket/export-id/export-id.overall_export_metadata",
+        "--async");
+
+    verify(datastoreAdmin, never()).get("opName");
+  }
+
+  @Test
+  public void test_failure_notAllowedInProduction() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () ->
+            runCommandInEnvironment(
+                RegistryToolEnvironment.PRODUCTION,
+                "--force",
+                "--poll_interval",
+                "PT0.001S",
+                "--backup_url",
+                "gs://bucket/export-id/export-id.overall_export_metadata"));
+  }
+
+  @Test
+  public void test_success_runInProduction() throws Exception {
+    runCommandInEnvironment(
+        RegistryToolEnvironment.PRODUCTION,
+        "--force",
+        "--confirm_production_import",
+        "PRODUCTION",
+        "--poll_interval",
+        "PT0.001S",
+        "--backup_url",
+        "gs://bucket/export-id/export-id.overall_export_metadata");
+  }
+}