Add TLDs parameter to refresh DNS action

This will allow us to migrate one TLD at a time by refreshing all zones
on the specified TLD after dual-writing is enabled.

Note that the TLDs parameter is required, which seems reasonable given
that almost all imagined use cases would be on a by-TLD basis.

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

java/google/registry/tools/server/RefreshDnsForAllDomainsAction.java

@@ -15,25 +15,31 @@
 package google.registry.tools.server;
 
 import static google.registry.mapreduce.inputs.EppResourceInputs.createEntityInput;
+import static google.registry.model.EppResourceUtils.isActive;
+import static google.registry.model.registry.Registries.assertTldsExist;
+import static google.registry.util.FormattingLogger.getLoggerForCallerClass;
 import static google.registry.util.PipelineUtils.createJobPath;
 
 import com.google.appengine.tools.mapreduce.Mapper;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import google.registry.dns.DnsQueue;
 import google.registry.mapreduce.MapreduceRunner;
-import google.registry.model.EppResourceUtils;
 import google.registry.model.domain.DomainResource;
 import google.registry.request.Action;
+import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.request.auth.AuthLevel;
 import google.registry.util.FormattingLogger;
+import google.registry.util.NonFinalForTesting;
 import javax.inject.Inject;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 
 /**
- * A mapreduce that enqueues DNS publish tasks on all active domains.
+ * A mapreduce that enqueues DNS publish tasks on all active domains on the specified TLD(s).
  *
  * <p>This refreshes DNS both for all domain names and all in-bailiwick hostnames, as DNS writers
  * are responsible for enqueuing refresh tasks for subordinate hosts. So this action thus refreshes
@@ -54,22 +60,24 @@ import org.joda.time.DateTimeZone;
 )
 public class RefreshDnsForAllDomainsAction implements Runnable {
 
-  private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
+  private static final FormattingLogger logger = getLoggerForCallerClass();
 
   @Inject MapreduceRunner mrRunner;
   @Inject Response response;
+  @Inject @Parameter("tlds") ImmutableSet<String> tlds;
   @Inject RefreshDnsForAllDomainsAction() {}
 
   @Override
   public void run() {
+    assertTldsExist(tlds);
     response.sendJavaScriptRedirect(
         createJobPath(
             mrRunner
-                .setJobName("Refresh all domains")
+                .setJobName("Refresh DNS for all domains")
                 .setModuleName("tools")
                 .setDefaultMapShards(10)
                 .runMapOnly(
-                    new RefreshDnsForAllDomainsActionMapper(),
+                    new RefreshDnsForAllDomainsActionMapper(tlds),
                     ImmutableList.of(createEntityInput(DomainResource.class)))));
   }
 
@@ -77,22 +85,35 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
   public static class RefreshDnsForAllDomainsActionMapper
       extends Mapper<DomainResource, Void, Void> {
 
-    private static final DnsQueue dnsQueue = DnsQueue.create();
-    private static final long serialVersionUID = 1356876487351666133L;
+    private static final long serialVersionUID = 1455544013508953083L;
+
+    @NonFinalForTesting
+    @VisibleForTesting
+    static DnsQueue dnsQueue = DnsQueue.create();
+
+    private final ImmutableSet<String> tlds;
+
+    RefreshDnsForAllDomainsActionMapper(ImmutableSet<String> tlds) {
+      this.tlds = tlds;
+    }
 
     @Override
     public void map(final DomainResource domain) {
       String domainName = domain.getFullyQualifiedDomainName();
-      if (EppResourceUtils.isActive(domain, DateTime.now(DateTimeZone.UTC))) {
-        try {
-          dnsQueue.addDomainRefreshTask(domainName);
-          getContext().incrementCounter("active domains refreshed");
-        } catch (Throwable t) {
-          logger.severefmt(t, "Error while refreshing DNS for domain %s", domainName);
-          getContext().incrementCounter("active domains errored");
+      if (tlds.contains(domain.getTld())) {
+        if (isActive(domain, DateTime.now(DateTimeZone.UTC))) {
+          try {
+            dnsQueue.addDomainRefreshTask(domainName);
+            getContext().incrementCounter("active domains refreshed");
+          } catch (Throwable t) {
+            logger.severefmt(t, "Error while refreshing DNS for domain %s", domainName);
+            getContext().incrementCounter("active domains errored");
+          }
+        } else {
+          getContext().incrementCounter("inactive domains skipped");
         }
       } else {
-        getContext().incrementCounter("inactive domains skipped");
+        getContext().incrementCounter("domains on non-targeted TLDs skipped");
       }
     }
   }

javatests/google/registry/tools/server/BUILD

@@ -12,6 +12,7 @@ java_library(
     srcs = glob(["*.java"]),
     resources = glob(["testdata/*"]),
     deps = [
+        "//java/google/registry/dns",
         "//java/google/registry/groups",
         "//java/google/registry/mapreduce",
         "//java/google/registry/model",

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

@@ -0,0 +1,96 @@
+// 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.server;
+
+import static google.registry.testing.DatastoreHelper.createTld;
+import static google.registry.testing.DatastoreHelper.createTlds;
+import static google.registry.testing.DatastoreHelper.persistActiveDomain;
+import static google.registry.testing.DatastoreHelper.persistDeletedDomain;
+import static org.joda.time.DateTimeZone.UTC;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import com.google.common.collect.ImmutableSet;
+import google.registry.dns.DnsQueue;
+import google.registry.testing.FakeResponse;
+import google.registry.testing.InjectRule;
+import google.registry.testing.mapreduce.MapreduceTestCase;
+import google.registry.tools.server.RefreshDnsForAllDomainsAction.RefreshDnsForAllDomainsActionMapper;
+import org.joda.time.DateTime;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link RefreshDnsForAllDomainsAction}. */
+@RunWith(JUnit4.class)
+public class RefreshDnsForAllDomainsActionTest
+    extends MapreduceTestCase<RefreshDnsForAllDomainsAction> {
+
+  @Rule public final InjectRule inject = new InjectRule();
+
+  private final DnsQueue dnsQueue = mock(DnsQueue.class);
+
+  @Before
+  public void init() {
+    inject.setStaticField(RefreshDnsForAllDomainsActionMapper.class, "dnsQueue", dnsQueue);
+
+    action = new RefreshDnsForAllDomainsAction();
+    action.mrRunner = makeDefaultRunner();
+    action.response = new FakeResponse();
+  }
+
+  private void runMapreduce() throws Exception {
+    action.run();
+    executeTasksUntilEmpty("mapreduce");
+  }
+
+  @Test
+  public void test_runAction_successfullyEnqueuesDnsRefreshes() throws Exception {
+    createTld("bar");
+    persistActiveDomain("foo.bar");
+    persistActiveDomain("low.bar");
+    action.tlds = ImmutableSet.of("bar");
+    runMapreduce();
+    verify(dnsQueue).addDomainRefreshTask("foo.bar");
+    verify(dnsQueue).addDomainRefreshTask("low.bar");
+  }
+
+  @Test
+  public void test_runAction_doesntRefreshDeletedDomain() throws Exception {
+    createTld("bar");
+    persistActiveDomain("foo.bar");
+    persistDeletedDomain("deleted.bar", DateTime.now(UTC).minusYears(1));
+    action.tlds = ImmutableSet.of("bar");
+    runMapreduce();
+    verify(dnsQueue).addDomainRefreshTask("foo.bar");
+    verify(dnsQueue, never()).addDomainRefreshTask("deleted.bar");
+  }
+
+  @Test
+  public void test_runAction_ignoresDomainsOnOtherTlds() throws Exception {
+    createTlds("bar", "baz");
+    persistActiveDomain("foo.bar");
+    persistActiveDomain("low.bar");
+    persistActiveDomain("ignore.baz");
+    action.tlds = ImmutableSet.of("bar");
+    runMapreduce();
+    verify(dnsQueue).addDomainRefreshTask("foo.bar");
+    verify(dnsQueue).addDomainRefreshTask("low.bar");
+    verify(dnsQueue, never()).addDomainRefreshTask("ignore.baz");
+  }
+}