Centralize creation of MapreduceRunner in unit tests

This also changes the default number of mapper shards in tests to 2, which is
the number of EppResourceIndex buckets in unit tests. Running more shards than
there are buckets causes unnecessary test load.

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

javatests/google/registry/batch/DeleteProberDataActionTest.java

@@ -23,10 +23,8 @@ import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.DatastoreHelper.persistSimpleResource;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 
-import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Reason;
@@ -74,7 +72,7 @@ public class DeleteProberDataActionTest extends MapreduceTestCase<DeleteProberDa
     persistResource(Registry.get("oa-canary.test").asBuilder().setTldType(TldType.TEST).build());
 
     action = new DeleteProberDataAction();
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>of(5), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
     action.isDryRun = false;
   }

javatests/google/registry/batch/ExpandRecurringBillingEventsActionTest.java

@@ -32,7 +32,6 @@ import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.VoidWork;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
@@ -71,7 +70,7 @@ public class ExpandRecurringBillingEventsActionTest
   @Before
   public void init() {
     action = new ExpandRecurringBillingEventsAction();
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>of(5), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.clock = clock;
     action.cursorTimeParam = Optional.absent();
     createTld("tld");

javatests/google/registry/export/ExportDomainListsActionTest.java

@@ -28,9 +28,7 @@ import com.google.appengine.tools.cloudstorage.GcsFilename;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.appengine.tools.cloudstorage.ListOptions;
 import com.google.appengine.tools.cloudstorage.ListResult;
-import com.google.common.base.Optional;
 import com.google.common.base.Splitter;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldType;
 import google.registry.testing.ExceptionRule;
@@ -60,7 +58,7 @@ public class ExportDomainListsActionTest extends MapreduceTestCase<ExportDomainL
     persistResource(Registry.get("testtld").asBuilder().setTldType(TldType.TEST).build());
 
     action = new ExportDomainListsAction();
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>absent(), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
     action.gcsBucket = "outputbucket";
     action.gcsBufferSize = 500;

javatests/google/registry/flows/async/DeleteContactsAndHostsActionTest.java

@@ -55,13 +55,11 @@ import static org.joda.time.Duration.standardSeconds;
 import com.google.appengine.api.taskqueue.QueueFactory;
 import com.google.appengine.api.taskqueue.TaskOptions;
 import com.google.appengine.api.taskqueue.TaskOptions.Method;
-import com.google.common.base.Optional;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.googlecode.objectify.Key;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.EppResource;
 import google.registry.model.contact.ContactAddress;
 import google.registry.model.contact.ContactPhoneNumber;
@@ -133,7 +131,7 @@ public class DeleteContactsAndHostsActionTest
 
     action = new DeleteContactsAndHostsAction();
     action.clock = clock;
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>of(5), Optional.<Integer>of(2));
+    action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
     action.retrier = new Retrier(new FakeSleeper(clock), 1);
     action.queue = getQueue(QUEUE_ASYNC_DELETE);

javatests/google/registry/flows/async/DnsRefreshForHostRenameActionTest.java

@@ -26,11 +26,9 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.verifyZeroInteractions;
 
-import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.dns.DnsQueue;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.host.HostResource;
 import google.registry.request.HttpException.BadRequestException;
 import google.registry.testing.ExceptionRule;
@@ -64,7 +62,7 @@ public class DnsRefreshForHostRenameActionTest
   private void runMapreduce(String hostKeyString) throws Exception {
     action = new DnsRefreshForHostRenameAction();
     action.hostKeyString = hostKeyString;
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>of(5), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
     action.run();
     executeTasksUntilEmpty("mapreduce");

javatests/google/registry/flows/async/RefreshDnsOnHostRenameActionTest.java

@@ -33,10 +33,8 @@ import static org.joda.time.Duration.millis;
 import static org.joda.time.Duration.standardHours;
 import static org.joda.time.Duration.standardSeconds;
 
-import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.host.HostResource;
 import google.registry.testing.ExceptionRule;
 import google.registry.testing.FakeClock;
@@ -79,7 +77,7 @@ public class RefreshDnsOnHostRenameActionTest
 
     action = new RefreshDnsOnHostRenameAction();
     action.clock = clock;
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>of(5), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.pullQueue = getQueue(QUEUE_ASYNC_HOST_RENAME);
     action.response = new FakeResponse();
     action.retrier = new Retrier(new FakeSleeper(clock), 1);

javatests/google/registry/rde/RdeContactImportActionTest.java

@@ -58,7 +58,7 @@ public class RdeContactImportActionTest extends MapreduceTestCase<RdeContactImpo
   @Before
   public void before() throws Exception {
     response = new FakeResponse();
-    mrRunner = new MapreduceRunner(Optional.<Integer>absent(), Optional.<Integer>absent());
+    mrRunner = makeDefaultRunner();
     action = new RdeContactImportAction(
         mrRunner,
         response,

javatests/google/registry/rde/RdeStagingActionTest.java

@@ -36,14 +36,12 @@ import static java.util.Arrays.asList;
 import com.google.appengine.tools.cloudstorage.GcsFilename;
 import com.google.appengine.tools.cloudstorage.GcsService;
 import com.google.appengine.tools.cloudstorage.GcsServiceFactory;
-import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
 import com.googlecode.objectify.VoidWork;
 import google.registry.keyring.api.Keyring;
 import google.registry.keyring.api.PgpHelper;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.common.Cursor;
 import google.registry.model.common.Cursor.CursorType;
 import google.registry.model.host.HostResource;
@@ -123,7 +121,7 @@ public class RdeStagingActionTest extends MapreduceTestCase<RdeStagingAction> {
     inject.setStaticField(Ofy.class, "clock", clock);
     action = new RdeStagingAction();
     action.clock = clock;
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>absent(), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.reducer = new RdeStagingReducer();
     action.reducer.ghostrydeBufferSize = 31337;
     action.reducer.lockTimeout = Duration.standardHours(1);

javatests/google/registry/testing/mapreduce/BUILD

@@ -26,6 +26,8 @@ java_library(
         "//third_party/java/mockito",
         "//third_party/java/servlet/servlet_api",
         "//third_party/java/truth",
+        "//java/google/registry/config",
+        "//java/google/registry/mapreduce",
         "//java/google/registry/model",
         "//java/google/registry/util",
         "//javatests/google/registry/testing",

javatests/google/registry/testing/mapreduce/MapreduceTestCase.java

@@ -32,6 +32,9 @@ import com.google.appengine.tools.pipeline.impl.servlets.PipelineServlet;
 import com.google.appengine.tools.pipeline.impl.servlets.TaskHandler;
 import com.google.apphosting.api.ApiProxy;
 import com.google.common.base.CharMatcher;
+import com.google.common.base.Optional;
+import google.registry.config.RegistryEnvironment;
+import google.registry.mapreduce.MapreduceRunner;
 import google.registry.testing.AppEngineRule;
 import google.registry.testing.FakeClock;
 import google.registry.testing.ShardableTestCase;
@@ -84,6 +87,11 @@ public abstract class MapreduceTestCase<T> extends ShardableTestCase {
     proxy.setProperty(LocalBlobstoreService.NO_STORAGE_PROPERTY, "true");
   }
 
+  protected MapreduceRunner makeDefaultRunner() {
+    int numBuckets = RegistryEnvironment.get().config().getEppResourceIndexBucketCount();
+    return new MapreduceRunner(Optional.<Integer>of(numBuckets), Optional.<Integer>of(1));
+  }
+
   protected List<QueueStateInfo.TaskStateInfo> getTasks(String queueName) {
     return taskQueue.getQueueStateInfo().get(queueName).getTaskInfo();
   }

javatests/google/registry/tools/server/GenerateZoneFilesActionTest.java

@@ -31,13 +31,11 @@ import static org.joda.time.Duration.standardDays;
 
 import com.google.appengine.tools.cloudstorage.GcsFilename;
 import com.google.appengine.tools.cloudstorage.GcsService;
-import com.google.common.base.Optional;
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.host.HostResource;
 import google.registry.testing.FakeClock;
@@ -97,7 +95,7 @@ public class GenerateZoneFilesActionTest extends MapreduceTestCase<GenerateZoneF
         .build());
 
     GenerateZoneFilesAction action = new GenerateZoneFilesAction();
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>absent(), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.bucket = "zonefiles-bucket";
     action.gcsBufferSize = 123;
     action.datastoreRetention = standardDays(29);

javatests/google/registry/tools/server/KillAllCommitLogsActionTest.java

@@ -27,12 +27,10 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static java.util.Arrays.asList;
 
 import com.google.appengine.api.datastore.Entity;
-import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.ImmutableObject;
 import google.registry.model.ofy.CommitLogBucket;
 import google.registry.model.ofy.CommitLogCheckpoint;
@@ -59,7 +57,7 @@ public class KillAllCommitLogsActionTest extends MapreduceTestCase<KillAllCommit
 
   private void runMapreduce() throws Exception {
     action = new KillAllCommitLogsAction();
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>absent(), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
     action.run();
     executeTasksUntilEmpty("mapreduce");

javatests/google/registry/tools/server/KillAllEppResourcesActionTest.java

@@ -32,12 +32,10 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static java.util.Arrays.asList;
 
 import com.google.appengine.api.datastore.Entity;
-import com.google.common.base.Optional;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.Iterables;
 import com.googlecode.objectify.Key;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
@@ -84,7 +82,7 @@ public class KillAllEppResourcesActionTest extends MapreduceTestCase<KillAllEppR
 
   private void runMapreduce() throws Exception {
     action = new KillAllEppResourcesAction();
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>absent(), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
     action.run();
     executeTasksUntilEmpty("mapreduce");

javatests/google/registry/tools/server/ResaveAllEppResourcesActionTest.java

@@ -18,8 +18,6 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.testing.DatastoreHelper.persistActiveContact;
 
-import com.google.common.base.Optional;
-import google.registry.mapreduce.MapreduceRunner;
 import google.registry.model.contact.ContactResource;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.mapreduce.MapreduceTestCase;
@@ -37,7 +35,7 @@ public class ResaveAllEppResourcesActionTest
   @Before
   public void init() {
     action = new ResaveAllEppResourcesAction();
-    action.mrRunner = new MapreduceRunner(Optional.<Integer>of(5), Optional.<Integer>absent());
+    action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
   }