Consolidate the use of URL parameters to specify database override (#1331)

There are actions for which we want to provide an override for the database
to use, like when launching Spec11 and Invoicing pipelines. It make sense to
consolidate around the same parameter provided from the same module for
consistency in all cases, instead of defining an override for each action.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/google/nomulus/1331)
<!-- Reviewable:end -->

core/src/main/java/google/registry/batch/BatchModule.java

@@ -24,6 +24,7 @@ import static google.registry.batch.AsyncTaskEnqueuer.QUEUE_ASYNC_HOST_RENAME;
 import static google.registry.request.RequestParameters.extractIntParameter;
 import static google.registry.request.RequestParameters.extractLongParameter;
 import static google.registry.request.RequestParameters.extractOptionalBooleanParameter;
+import static google.registry.request.RequestParameters.extractOptionalDatetimeParameter;
 import static google.registry.request.RequestParameters.extractOptionalIntParameter;
 import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredDatetimeParameter;
@@ -106,6 +107,13 @@ public class BatchModule {
     return extractIntParameter(req, RelockDomainAction.PREVIOUS_ATTEMPTS_PARAM);
   }
 
+  @Provides
+  @Parameter(ExpandRecurringBillingEventsAction.PARAM_CURSOR_TIME)
+  static Optional<DateTime> provideCursorTime(HttpServletRequest req) {
+    return extractOptionalDatetimeParameter(
+        req, ExpandRecurringBillingEventsAction.PARAM_CURSOR_TIME);
+  }
+
   @Provides
   @Named(QUEUE_ASYNC_ACTIONS)
   static Queue provideAsyncActionsPushQueue() {

core/src/main/java/google/registry/module/backend/BackendModule.java

@@ -1,59 +0,0 @@
-// 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.module.backend;
-
-import static google.registry.model.tld.Registries.assertTldExists;
-import static google.registry.model.tld.Registries.assertTldsExist;
-import static google.registry.request.RequestParameters.extractOptionalDatetimeParameter;
-import static google.registry.request.RequestParameters.extractRequiredParameter;
-import static google.registry.request.RequestParameters.extractSetOfParameters;
-
-import com.google.common.collect.ImmutableSet;
-import dagger.Module;
-import dagger.Provides;
-import google.registry.batch.ExpandRecurringBillingEventsAction;
-import google.registry.request.Parameter;
-import google.registry.request.RequestParameters;
-import java.util.Optional;
-import javax.servlet.http.HttpServletRequest;
-import org.joda.time.DateTime;
-
-/**
- * Dagger module for injecting common settings for all Backend tasks.
- */
-@Module
-public class BackendModule {
-
-  @Provides
-  @Parameter(RequestParameters.PARAM_TLD)
-  static String provideTld(HttpServletRequest req) {
-    return assertTldExists(extractRequiredParameter(req, RequestParameters.PARAM_TLD));
-  }
-
-  @Provides
-  @Parameter(RequestParameters.PARAM_TLDS)
-  static ImmutableSet<String> provideTlds(HttpServletRequest req) {
-    ImmutableSet<String> tlds = extractSetOfParameters(req, RequestParameters.PARAM_TLDS);
-    assertTldsExist(tlds);
-    return tlds;
-  }
-
-  @Provides
-  @Parameter("cursorTime")
-  static Optional<DateTime> provideCursorTime(HttpServletRequest req) {
-    return extractOptionalDatetimeParameter(
-        req, ExpandRecurringBillingEventsAction.PARAM_CURSOR_TIME);
-  }
-}

core/src/main/java/google/registry/module/backend/BackendRequestComponent.java

@@ -94,7 +94,6 @@ import google.registry.tools.javascrap.CreateSyntheticHistoryEntriesAction;
 @RequestScope
 @Subcomponent(
     modules = {
-      BackendModule.class,
       BackupModule.class,
       BatchModule.class,
       BillingModule.class,

core/src/main/java/google/registry/reporting/ReportingModule.java

@@ -14,7 +14,6 @@
 
 package google.registry.reporting;
 
-import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 
@@ -56,9 +55,6 @@ public class ReportingModule {
   /** The request parameter specifying the jobId for a running Dataflow pipeline. */
   public static final String PARAM_JOB_ID = "jobId";
 
-  /** The request parameter for specifying which database reporting actions should read from. */
-  public static final String DATABASE = "database";
-
   /** Provides the Cloud Dataflow jobId for a pipeline. */
   @Provides
   @Parameter(PARAM_JOB_ID)
@@ -66,14 +62,6 @@ public class ReportingModule {
     return extractRequiredParameter(req, PARAM_JOB_ID);
   }
 
-  /** Provides the database for the pipeline to read from. */
-  @Provides
-  @Parameter(DATABASE)
-  static String provideDatabase(HttpServletRequest req) {
-    Optional<String> optionalDatabase = extractOptionalParameter(req, DATABASE);
-    return optionalDatabase.orElse(tm().isOfy() ? "DATASTORE" : "CLOUD_SQL");
-  }
-
   /** Extracts an optional YearMonth in yyyy-MM format from the request. */
   @Provides
   @Parameter(PARAM_YEAR_MONTH)

core/src/main/java/google/registry/reporting/billing/GenerateInvoicesAction.java

@@ -15,10 +15,9 @@
 package google.registry.reporting.billing;
 
 import static google.registry.beam.BeamUtils.createJobName;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase.CLOUD_SQL;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.reporting.ReportingModule.DATABASE;
 import static google.registry.reporting.ReportingUtils.enqueueBeamReportingTask;
-import static google.registry.reporting.billing.BillingModule.PARAM_SHOULD_PUBLISH;
 import static google.registry.request.Action.Method.POST;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
@@ -31,9 +30,11 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.net.MediaType;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
 import google.registry.reporting.ReportingModule;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
+import google.registry.request.RequestParameters;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.util.Clock;
@@ -72,7 +73,7 @@ public class GenerateInvoicesAction implements Runnable {
   private final Clock clock;
   private final Response response;
   private final Dataflow dataflow;
-  private final String database;
+  private final PrimaryDatabase database;
 
   @Inject
   GenerateInvoicesAction(
@@ -81,8 +82,8 @@ public class GenerateInvoicesAction implements Runnable {
       @Config("beamStagingBucketUrl") String stagingBucketUrl,
       @Config("billingBucketUrl") String billingBucketUrl,
       @Config("invoiceFilePrefix") String invoiceFilePrefix,
-      @Parameter(PARAM_SHOULD_PUBLISH) boolean shouldPublish,
+      @Parameter(BillingModule.PARAM_SHOULD_PUBLISH) boolean shouldPublish,
-      @Parameter(DATABASE) String database,
+      @Parameter(RequestParameters.PARAM_DATABASE) PrimaryDatabase database,
       YearMonth yearMonth,
       BillingEmailUtils emailUtils,
       Clock clock,
@@ -93,7 +94,7 @@ public class GenerateInvoicesAction implements Runnable {
     this.stagingBucketUrl = stagingBucketUrl;
     // When generating the invoices using Cloud SQL before database cutover, save the reports in a
     // separate bucket so that it does not overwrite the Datastore invoices.
-    if (tm().isOfy() && database.equals("CLOUD_SQL")) {
+    if (tm().isOfy() && database.equals(CLOUD_SQL)) {
       billingBucketUrl = billingBucketUrl.concat("-sql");
     }
     this.billingBucketUrl = billingBucketUrl;
@@ -124,7 +125,7 @@ public class GenerateInvoicesAction implements Runnable {
                       "invoiceFilePrefix",
                       invoiceFilePrefix,
                       "database",
-                      database,
+                      database.name(),
                       "billingBucketUrl",
                       billingBucketUrl));
       LaunchFlexTemplateResponse launchResponse =
@@ -148,14 +149,13 @@ public class GenerateInvoicesAction implements Runnable {
                 yearMonth.toString());
         enqueueBeamReportingTask(PublishInvoicesAction.PATH, beamTaskParameters);
       }
+      response.setStatus(SC_OK);
+      response.setPayload(String.format("Launched invoicing pipeline: %s", jobId));
     } catch (IOException e) {
-      logger.atWarning().withCause(e).log("Template Launch failed");
+      logger.atWarning().withCause(e).log("Pipeline Launch failed");
-      emailUtils.sendAlertEmail(String.format("Template Launch failed due to %s", e.getMessage()));
+      emailUtils.sendAlertEmail(String.format("Pipeline Launch failed due to %s", e.getMessage()));
       response.setStatus(SC_INTERNAL_SERVER_ERROR);
-      response.setPayload(String.format("Template launch failed: %s", e.getMessage()));
+      response.setPayload(String.format("Pipeline launch failed: %s", e.getMessage()));
-      return;
     }
-    response.setStatus(SC_OK);
-    response.setPayload("Launched dataflow template.");
   }
 }

core/src/main/java/google/registry/reporting/spec11/GenerateSpec11ReportAction.java

@@ -16,8 +16,6 @@ package google.registry.reporting.spec11;
 
 import static google.registry.beam.BeamUtils.createJobName;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
-import static google.registry.reporting.ReportingModule.DATABASE;
-import static google.registry.reporting.ReportingModule.PARAM_DATE;
 import static google.registry.reporting.ReportingUtils.enqueueBeamReportingTask;
 import static google.registry.request.Action.Method.POST;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
@@ -33,9 +31,11 @@ import com.google.common.net.MediaType;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.config.RegistryEnvironment;
 import google.registry.keyring.api.KeyModule.Key;
+import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
 import google.registry.reporting.ReportingModule;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
+import google.registry.request.RequestParameters;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.util.Clock;
@@ -71,7 +71,7 @@ public class GenerateSpec11ReportAction implements Runnable {
   private final Clock clock;
   private final Response response;
   private final Dataflow dataflow;
-  private final String database;
+  private final PrimaryDatabase database;
 
   @Inject
   GenerateSpec11ReportAction(
@@ -80,15 +80,15 @@ public class GenerateSpec11ReportAction implements Runnable {
       @Config("beamStagingBucketUrl") String stagingBucketUrl,
       @Config("reportingBucketUrl") String reportingBucketUrl,
       @Key("safeBrowsingAPIKey") String apiKey,
-      @Parameter(PARAM_DATE) LocalDate date,
+      @Parameter(ReportingModule.PARAM_DATE) LocalDate date,
-      @Parameter(DATABASE) String database,
+      @Parameter(RequestParameters.PARAM_DATABASE) PrimaryDatabase database,
       Clock clock,
       Response response,
       Dataflow dataflow) {
     this.projectId = projectId;
     this.jobRegion = jobRegion;
     this.stagingBucketUrl = stagingBucketUrl;
-    if (tm().isOfy() && database.equals("CLOUD_SQL")) {
+    if (tm().isOfy() && database.equals(PrimaryDatabase.CLOUD_SQL)) {
       reportingBucketUrl = reportingBucketUrl.concat("-sql");
     }
     this.reportingBucketUrl = reportingBucketUrl;
@@ -114,7 +114,7 @@ public class GenerateSpec11ReportAction implements Runnable {
                       "safeBrowsingApiKey",
                       apiKey,
                       "database",
-                      database,
+                      database.name(),
                       ReportingModule.PARAM_DATE,
                       date.toString(),
                       "reportingBucketUrl",
@@ -131,21 +131,18 @@ public class GenerateSpec11ReportAction implements Runnable {
                   jobRegion,
                   new LaunchFlexTemplateRequest().setLaunchParameter(parameter))
               .execute();
+      logger.atInfo().log("Got response: %s", launchResponse.getJob().toPrettyString());
+      String jobId = launchResponse.getJob().getId();
       Map<String, String> beamTaskParameters =
           ImmutableMap.of(
-              ReportingModule.PARAM_JOB_ID,
+              ReportingModule.PARAM_JOB_ID, jobId, ReportingModule.PARAM_DATE, date.toString());
-              launchResponse.getJob().getId(),
-              ReportingModule.PARAM_DATE,
-              date.toString());
       enqueueBeamReportingTask(PublishSpec11ReportAction.PATH, beamTaskParameters);
-      logger.atInfo().log("Got response: %s", launchResponse.getJob().toPrettyString());
+      response.setStatus(SC_OK);
+      response.setPayload(String.format("Launched Spec11 pipeline: %s", jobId));
     } catch (IOException e) {
-      logger.atWarning().withCause(e).log("Template Launch failed");
+      logger.atWarning().withCause(e).log("Pipeline Launch failed");
       response.setStatus(SC_INTERNAL_SERVER_ERROR);
-      response.setPayload(String.format("Template launch failed: %s", e.getMessage()));
+      response.setPayload(String.format("Pipeline launch failed: %s", e.getMessage()));
-      return;
     }
-    response.setStatus(SC_OK);
-    response.setPayload("Launched Spec11 dataflow template.");
   }
 }

core/src/main/java/google/registry/request/RequestModule.java

@@ -15,16 +15,24 @@
 package google.registry.request;
 
 import static com.google.common.net.MediaType.JSON_UTF_8;
+import static google.registry.model.tld.Registries.assertTldExists;
+import static google.registry.model.tld.Registries.assertTldsExist;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.request.RequestParameters.extractOptionalParameter;
+import static google.registry.request.RequestParameters.extractRequiredParameter;
+import static google.registry.request.RequestParameters.extractSetOfParameters;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.VerifyException;
 import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.io.ByteStreams;
 import com.google.common.io.CharStreams;
 import com.google.common.net.MediaType;
 import dagger.Module;
 import dagger.Provides;
+import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
 import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.HttpException.UnsupportedMediaTypeException;
 import google.registry.request.auth.AuthResult;
@@ -60,6 +68,29 @@ public final class RequestModule {
     this.authResult = authResult;
   }
 
+  @Provides
+  @Parameter(RequestParameters.PARAM_TLD)
+  static String provideTld(HttpServletRequest req) {
+    return assertTldExists(extractRequiredParameter(req, RequestParameters.PARAM_TLD));
+  }
+
+  @Provides
+  @Parameter(RequestParameters.PARAM_TLDS)
+  static ImmutableSet<String> provideTlds(HttpServletRequest req) {
+    ImmutableSet<String> tlds = extractSetOfParameters(req, RequestParameters.PARAM_TLDS);
+    assertTldsExist(tlds);
+    return tlds;
+  }
+
+  @Provides
+  @Parameter(RequestParameters.PARAM_DATABASE)
+  static PrimaryDatabase provideDatabase(HttpServletRequest req) {
+    return extractOptionalParameter(req, RequestParameters.PARAM_DATABASE)
+        .map(String::toUpperCase)
+        .map(PrimaryDatabase::valueOf)
+        .orElse(tm().isOfy() ? PrimaryDatabase.DATASTORE : PrimaryDatabase.CLOUD_SQL);
+  }
+
   @Provides
   static Response provideResponse(ResponseImpl response) {
     return response;

core/src/main/java/google/registry/request/RequestParameters.java

@@ -37,6 +37,14 @@ public final class RequestParameters {
   /** The standardized request parameter name used by any action taking multiple tld parameters. */
   public static final String PARAM_TLDS = "tlds";
 
+  /**
+   * The standardized optional request parameter name used by any action to specify which database
+   * to use, if the action supports such override. The supported values are (case-insensitive)
+   * "datastore" and "cloud_sql".
+   */
+  // TODO (jianglai): delete this param after the database migration.
+  public static final String PARAM_DATABASE = "database";
+
   /**
    * Returns first GET or POST parameter associated with {@code name}.
    *

core/src/main/java/google/registry/tools/server/ToolsServerModule.java

@@ -18,13 +18,10 @@ import static com.google.common.base.Strings.emptyToNull;
 import static google.registry.request.RequestParameters.extractIntParameter;
 import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
-import static google.registry.request.RequestParameters.extractSetOfParameters;
 
-import com.google.common.collect.ImmutableSet;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.request.Parameter;
-import google.registry.request.RequestParameters;
 import java.util.Optional;
 import javax.servlet.http.HttpServletRequest;
 
@@ -60,18 +57,6 @@ public class ToolsServerModule {
     return (s == null) ? Optional.empty() : Optional.of(Boolean.parseBoolean(s));
   }
 
-  @Provides
-  @Parameter(RequestParameters.PARAM_TLD)
-  static String provideTld(HttpServletRequest req) {
-    return extractRequiredParameter(req, RequestParameters.PARAM_TLD);
-  }
-
-  @Provides
-  @Parameter(RequestParameters.PARAM_TLDS)
-  static ImmutableSet<String> provideTlds(HttpServletRequest req) {
-    return extractSetOfParameters(req, RequestParameters.PARAM_TLDS);
-  }
-
   @Provides
   @Parameter("limit")
   static int provideLimit(HttpServletRequest req) {

core/src/test/java/google/registry/dns/DnsTestComponent.java

@@ -18,7 +18,6 @@ import dagger.Component;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.cron.CronModule;
 import google.registry.dns.writer.VoidDnsWriterModule;
-import google.registry.module.backend.BackendModule;
 import google.registry.request.RequestModule;
 import google.registry.util.UtilsModule;
 import javax.inject.Singleton;
@@ -26,7 +25,6 @@ import javax.inject.Singleton;
 @Singleton
 @Component(
     modules = {
-      BackendModule.class,
       ConfigModule.class,
       CronModule.class,
       DnsModule.class,

core/src/test/java/google/registry/reporting/billing/GenerateInvoicesActionTest.java

@@ -25,15 +25,18 @@ import static org.mockito.Mockito.when;
 
 import com.google.common.net.MediaType;
 import google.registry.beam.BeamActionTestBase;
+import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DualDatabaseTest;
 import google.registry.testing.FakeClock;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
+import google.registry.testing.TestOfyAndSql;
 import java.io.IOException;
 import org.joda.time.YearMonth;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link google.registry.reporting.billing.GenerateInvoicesAction}. */
+@DualDatabaseTest
 class GenerateInvoicesActionTest extends BeamActionTestBase {
 
   @RegisterExtension
@@ -44,7 +47,7 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
   private FakeClock clock = new FakeClock();
   private GenerateInvoicesAction action;
 
-  @Test
+  @TestOfyAndSql
   void testLaunchTemplateJob_withPublish() throws Exception {
     action =
         new GenerateInvoicesAction(
@@ -54,7 +57,7 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
             "billing_bucket",
             "REG-INV",
             true,
-            "DATASTORE",
+            PrimaryDatabase.DATASTORE,
             new YearMonth(2017, 10),
             emailUtils,
             clock,
@@ -63,7 +66,7 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
     action.run();
     assertThat(response.getContentType()).isEqualTo(MediaType.PLAIN_TEXT_UTF_8);
     assertThat(response.getStatus()).isEqualTo(SC_OK);
-    assertThat(response.getPayload()).isEqualTo("Launched dataflow template.");
+    assertThat(response.getPayload()).isEqualTo("Launched invoicing pipeline: jobid");
 
     TaskMatcher matcher =
         new TaskMatcher()
@@ -74,7 +77,7 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
     assertTasksEnqueued("beam-reporting", matcher);
   }
 
-  @Test
+  @TestOfyAndSql
   void testLaunchTemplateJob_withoutPublish() throws Exception {
     action =
         new GenerateInvoicesAction(
@@ -84,7 +87,7 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
             "billing_bucket",
             "REG-INV",
             false,
-            "DATASTORE",
+            PrimaryDatabase.DATASTORE,
             new YearMonth(2017, 10),
             emailUtils,
             clock,
@@ -93,11 +96,11 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
     action.run();
     assertThat(response.getContentType()).isEqualTo(MediaType.PLAIN_TEXT_UTF_8);
     assertThat(response.getStatus()).isEqualTo(SC_OK);
-    assertThat(response.getPayload()).isEqualTo("Launched dataflow template.");
+    assertThat(response.getPayload()).isEqualTo("Launched invoicing pipeline: jobid");
     assertNoTasksEnqueued("beam-reporting");
   }
 
-  @Test
+  @TestOfyAndSql
   void testCaughtIOException() throws IOException {
     when(launch.execute()).thenThrow(new IOException("Pipeline error"));
     action =
@@ -108,7 +111,7 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
             "billing_bucket",
             "REG-INV",
             false,
-            "DATASTORE",
+            PrimaryDatabase.DATASTORE,
             new YearMonth(2017, 10),
             emailUtils,
             clock,
@@ -116,8 +119,8 @@ class GenerateInvoicesActionTest extends BeamActionTestBase {
             dataflow);
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_INTERNAL_SERVER_ERROR);
-    assertThat(response.getPayload()).isEqualTo("Template launch failed: Pipeline error");
+    assertThat(response.getPayload()).isEqualTo("Pipeline launch failed: Pipeline error");
-    verify(emailUtils).sendAlertEmail("Template Launch failed due to Pipeline error");
+    verify(emailUtils).sendAlertEmail("Pipeline Launch failed due to Pipeline error");
     assertNoTasksEnqueued("beam-reporting");
   }
 }

core/src/test/java/google/registry/reporting/spec11/GenerateSpec11ReportActionTest.java

@@ -23,6 +23,7 @@ import static org.mockito.Mockito.when;
 
 import com.google.common.net.MediaType;
 import google.registry.beam.BeamActionTestBase;
+import google.registry.model.common.DatabaseMigrationStateSchedule.PrimaryDatabase;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.FakeClock;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
@@ -51,7 +52,7 @@ class GenerateSpec11ReportActionTest extends BeamActionTestBase {
             "gs://reporting-project/reporting-bucket/",
             "api_key/a",
             clock.nowUtc().toLocalDate(),
-            "DATASTORE",
+            PrimaryDatabase.DATASTORE,
             clock,
             response,
             dataflow);
@@ -73,13 +74,14 @@ class GenerateSpec11ReportActionTest extends BeamActionTestBase {
             "gs://reporting-project/reporting-bucket/",
             "api_key/a",
             clock.nowUtc().toLocalDate(),
-            "DATASTORE",
+            PrimaryDatabase.DATASTORE,
             clock,
             response,
             dataflow);
     action.run();
     assertThat(response.getStatus()).isEqualTo(SC_OK);
     assertThat(response.getContentType()).isEqualTo(MediaType.PLAIN_TEXT_UTF_8);
+    assertThat(response.getPayload()).isEqualTo("Launched Spec11 pipeline: jobid");
     TaskMatcher matcher =
         new TaskMatcher()
             .url("/_dr/task/publishSpec11")