Define TestRule that manages cache use in tests

All current tests that use caches with custom data expiry values
now restore the default config when teardown. We need to prevent
new unsafe uses from being introduced.

Restoration code have also been added to a few other tests that modifies
static fields.

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

java/google/registry/model/EppResource.java

@@ -50,6 +50,7 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import org.joda.time.DateTime;
+import org.joda.time.Duration;
 
 /** An EPP entity object (i.e. a domain, contact, or host). */
 public abstract class EppResource extends BackupGroupRoot implements Buildable {
@@ -343,14 +344,20 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
    */
   @NonFinalForTesting
   private static LoadingCache<Key<? extends EppResource>, EppResource> cacheEppResources =
-      CacheBuilder.newBuilder()
+      createEppResourcesCache(getEppResourceCachingDuration());
-          .expireAfterWrite(getEppResourceCachingDuration().getMillis(), MILLISECONDS)
+
-          .maximumSize(getEppResourceMaxCachedEntries())
+  private static LoadingCache<Key<? extends EppResource>, EppResource> createEppResourcesCache(
-          .build(CACHE_LOADER);
+      Duration expiry) {
+    return CacheBuilder.newBuilder()
+        .expireAfterWrite(expiry.getMillis(), MILLISECONDS)
+        .maximumSize(getEppResourceMaxCachedEntries())
+        .build(CACHE_LOADER);
+  }
 
   @VisibleForTesting
-  public static void setCacheForTest(CacheBuilder<Object, Object> cacheBuilder) {
+  public static void setCacheForTest(Optional<Duration> expiry) {
-    cacheEppResources = cacheBuilder.build(CACHE_LOADER);
+    Duration effectiveExpiry = expiry.orElse(getEppResourceCachingDuration());
+    cacheEppResources = createEppResourcesCache(effectiveExpiry);
   }
 
   private static ImmutableMap<Key<? extends EppResource>, EppResource> loadMultiple(

java/google/registry/model/index/ForeignKeyIndex.java

@@ -48,6 +48,7 @@ import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nullable;
 import org.joda.time.DateTime;
+import org.joda.time.Duration;
 
 /**
  * Class to map a foreign key to the active instance of {@link EppResource} whose unique id matches
@@ -220,15 +221,20 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    */
   @NonFinalForTesting
   private static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
-      cacheForeignKeyIndexes =
+      cacheForeignKeyIndexes = createForeignKeyIndexesCache(getEppResourceCachingDuration());
-          CacheBuilder.newBuilder()
+
-              .expireAfterWrite(getEppResourceCachingDuration().getMillis(), MILLISECONDS)
+  private static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
-              .maximumSize(getEppResourceMaxCachedEntries())
+      createForeignKeyIndexesCache(Duration expiry) {
-              .build(CACHE_LOADER);
+    return CacheBuilder.newBuilder()
+        .expireAfterWrite(expiry.getMillis(), MILLISECONDS)
+        .maximumSize(getEppResourceMaxCachedEntries())
+        .build(CACHE_LOADER);
+  }
 
   @VisibleForTesting
-  public static void setCacheForTest(CacheBuilder<Object, Object> cacheBuilder) {
+  public static void setCacheForTest(Optional<Duration> expiry) {
-    cacheForeignKeyIndexes = cacheBuilder.build(CACHE_LOADER);
+    Duration effectiveExpiry = expiry.orElse(getEppResourceCachingDuration());
+    cacheForeignKeyIndexes = createForeignKeyIndexesCache(effectiveExpiry);
   }
 
   /**

java/google/registry/model/registry/label/PremiumList.java

@@ -136,6 +136,12 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
   static LoadingCache<String, PremiumList> cachePremiumLists =
       createCachePremiumLists(getDomainLabelListCacheDuration());
 
+  @VisibleForTesting
+  public static void setPremiumListCacheForTest(Optional<Duration> expiry) {
+    Duration effectiveExpiry = expiry.orElse(getDomainLabelListCacheDuration());
+    cachePremiumLists = createCachePremiumLists(effectiveExpiry);
+  }
+
   @VisibleForTesting
   static LoadingCache<String, PremiumList> createCachePremiumLists(Duration cachePersistDuration) {
     return CacheBuilder.newBuilder()
@@ -188,10 +194,16 @@ public final class PremiumList extends BaseDomainLabelList<Money, PremiumList.Pr
    * that exist, as well as those that might exist according to the Bloom filter, must be cached).
    * The entries judged least likely to be accessed again will be evicted first.
    */
-  @NonFinalForTesting @VisibleForTesting
+  @NonFinalForTesting
   static LoadingCache<Key<PremiumListEntry>, Optional<PremiumListEntry>> cachePremiumListEntries =
       createCachePremiumListEntries(getSingletonCachePersistDuration());
 
+  @VisibleForTesting
+  public static void setPremiumListEntriesCacheForTest(Optional<Duration> expiry) {
+    Duration effectiveExpiry = expiry.orElse(getSingletonCachePersistDuration());
+    cachePremiumListEntries = createCachePremiumListEntries(effectiveExpiry);
+  }
+
   @VisibleForTesting
   static LoadingCache<Key<PremiumListEntry>, Optional<PremiumListEntry>>
       createCachePremiumListEntries(Duration cachePersistDuration) {

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

@@ -79,9 +79,7 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
 
     private static final long serialVersionUID = 1455544013508953083L;
 
-    @NonFinalForTesting
+    @NonFinalForTesting private static DnsQueue dnsQueue = DnsQueue.create();
-    @VisibleForTesting
-    static DnsQueue dnsQueue = DnsQueue.create();
 
     private final ImmutableSet<String> tlds;
 
@@ -109,5 +107,12 @@ public class RefreshDnsForAllDomainsAction implements Runnable {
         getContext().incrementCounter("domains on non-targeted TLDs skipped");
       }
     }
+
+    @VisibleForTesting
+    public static DnsQueue setDnsQueueForTest(DnsQueue testQueue) {
+      DnsQueue currentQueue = dnsQueue;
+      dnsQueue = testQueue;
+      return currentQueue;
+    }
   }
 }

javatests/google/registry/model/EppResourceTest.java

@@ -20,21 +20,25 @@ import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.testing.DatastoreHelper.persistActiveContact;
 import static google.registry.testing.DatastoreHelper.persistActiveHost;
 import static google.registry.testing.DatastoreHelper.persistResource;
-import static java.util.concurrent.TimeUnit.DAYS;
 
-import com.google.common.cache.CacheBuilder;
 import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.host.HostResource;
+import google.registry.testing.TestCacheRule;
+import org.joda.time.Duration;
+import org.junit.Rule;
 import org.junit.Test;
 
 /** Unit tests for {@link EppResource}. */
 public class EppResourceTest extends EntityTestCase {
 
+  @Rule
+  public final TestCacheRule testCacheRule =
+      new TestCacheRule.Builder().withEppResourceCache(Duration.standardDays(1)).build();
+
   @Test
   public void test_loadCached_ignoresContactChange() {
-    setNonZeroCachingInterval();
     ContactResource originalContact = persistActiveContact("contact123");
     assertThat(EppResource.loadCached(ImmutableList.of(Key.create(originalContact))))
         .containsExactly(Key.create(originalContact), originalContact);
@@ -48,7 +52,6 @@ public class EppResourceTest extends EntityTestCase {
 
   @Test
   public void test_loadCached_ignoresHostChange() {
-    setNonZeroCachingInterval();
     HostResource originalHost = persistActiveHost("ns1.example.com");
     assertThat(EppResource.loadCached(ImmutableList.of(Key.create(originalHost))))
         .containsExactly(Key.create(originalHost), originalHost);
@@ -60,8 +63,4 @@ public class EppResourceTest extends EntityTestCase {
     assertThat(loadByForeignKey(HostResource.class, "ns1.example.com", clock.nowUtc()))
         .hasValue(modifiedHost);
   }
-
-  private static void setNonZeroCachingInterval() {
-    EppResource.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
-  }
 }

javatests/google/registry/model/index/ForeignKeyIndexTest.java

@@ -25,21 +25,26 @@ import static google.registry.testing.DatastoreHelper.persistActiveHost;
 import static google.registry.testing.DatastoreHelper.persistDeletedHost;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
-import static java.util.concurrent.TimeUnit.DAYS;
 
-import com.google.common.cache.CacheBuilder;
 import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import google.registry.model.EntityTestCase;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.host.HostResource;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
+import google.registry.testing.TestCacheRule;
+import org.joda.time.Duration;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 
 /** Unit tests for {@link ForeignKeyIndex}. */
 public class ForeignKeyIndexTest extends EntityTestCase {
 
+  @Rule
+  public final TestCacheRule testCacheRule =
+      new TestCacheRule.Builder().withForeignIndexKeyCache(Duration.standardDays(1)).build();
+
   @Before
   public void setUp() {
     createTld("com");
@@ -120,7 +125,6 @@ public class ForeignKeyIndexTest extends EntityTestCase {
 
   @Test
   public void test_loadCached_cachesNonexistenceOfHosts() {
-    setNonZeroCachingInterval();
     assertThat(
             ForeignKeyIndex.loadCached(
                 HostResource.class,
@@ -141,7 +145,6 @@ public class ForeignKeyIndexTest extends EntityTestCase {
 
   @Test
   public void test_loadCached_cachesExistenceOfHosts() {
-    setNonZeroCachingInterval();
     HostResource host1 = persistActiveHost("ns1.example.com");
     HostResource host2 = persistActiveHost("ns2.example.com");
     assertThat(
@@ -149,8 +152,11 @@ public class ForeignKeyIndexTest extends EntityTestCase {
                 HostResource.class,
                 ImmutableList.of("ns1.example.com", "ns2.example.com"),
                 clock.nowUtc()))
-        .containsExactly("ns1.example.com", loadHostFki("ns1.example.com"),
+        .containsExactly(
-            "ns2.example.com", loadHostFki("ns2.example.com"));
+            "ns1.example.com",
+            loadHostFki("ns1.example.com"),
+            "ns2.example.com",
+            loadHostFki("ns2.example.com"));
     deleteResource(host1);
     deleteResource(host2);
     persistActiveHost("ns3.example.com");
@@ -167,7 +173,6 @@ public class ForeignKeyIndexTest extends EntityTestCase {
 
   @Test
   public void test_loadCached_doesntSeeHostChangesWhileCacheIsValid() {
-    setNonZeroCachingInterval();
     HostResource originalHost = persistActiveHost("ns1.example.com");
     ForeignKeyIndex<HostResource> originalFki = loadHostFki("ns1.example.com");
     clock.advanceOneMilli();
@@ -191,7 +196,6 @@ public class ForeignKeyIndexTest extends EntityTestCase {
 
   @Test
   public void test_loadCached_filtersOutSoftDeletedHosts() {
-    setNonZeroCachingInterval();
     persistActiveHost("ns1.example.com");
     persistDeletedHost("ns2.example.com", clock.nowUtc().minusDays(1));
     assertThat(
@@ -204,7 +208,6 @@ public class ForeignKeyIndexTest extends EntityTestCase {
 
   @Test
   public void test_loadCached_cachesContactFkis() {
-    setNonZeroCachingInterval();
     persistActiveContact("contactid1");
     ForeignKeyIndex<ContactResource> fki1 = loadContactFki("contactid1");
     assertThat(
@@ -228,8 +231,4 @@ public class ForeignKeyIndexTest extends EntityTestCase {
             clock.nowUtc()))
         .containsExactly("contactid2", loadContactFki("contactid2"));
   }
-
-  private static void setNonZeroCachingInterval() {
-    ForeignKeyIndex.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
-  }
 }

javatests/google/registry/model/registry/label/PremiumListUtilsTest.java

@@ -26,7 +26,6 @@ import static google.registry.model.registry.label.DomainLabelMetrics.PremiumLis
 import static google.registry.model.registry.label.DomainLabelMetrics.PremiumListCheckOutcome.UNCACHED_POSITIVE;
 import static google.registry.model.registry.label.DomainLabelMetrics.premiumListChecks;
 import static google.registry.model.registry.label.DomainLabelMetrics.premiumListProcessingTime;
-import static google.registry.model.registry.label.PremiumList.createCachePremiumLists;
 import static google.registry.model.registry.label.PremiumListUtils.deletePremiumList;
 import static google.registry.model.registry.label.PremiumListUtils.doesPremiumListExist;
 import static google.registry.model.registry.label.PremiumListUtils.getPremiumPrice;
@@ -48,6 +47,7 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.label.PremiumList.PremiumListEntry;
 import google.registry.model.registry.label.PremiumList.PremiumListRevision;
 import google.registry.testing.AppEngineRule;
+import google.registry.testing.TestCacheRule;
 import java.util.Map;
 import org.joda.money.Money;
 import org.junit.Before;
@@ -62,13 +62,17 @@ public class PremiumListUtilsTest {
 
   @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
 
+  // Set long persist times on caches so they can be tested (cache times default to 0 in tests).
+  @Rule
+  public final TestCacheRule testCacheRule =
+      new TestCacheRule.Builder()
+          .withPremiumListsCache(standardDays(1))
+          .withPremiumListEntriesCache(standardDays(1))
+          .build();
+
   @Before
   public void before() {
-    // Set long persist times on caches so they can be tested (cache times default to 0 in tests).
+    // createTld() overwrites the premium list, so call it before persisting pl.
-    PremiumList.cachePremiumListEntries =
-        PremiumList.createCachePremiumListEntries(standardDays(1));
-    PremiumList.cachePremiumLists = createCachePremiumLists(standardDays(1));
-    // createTld() overwrites the premium list, so call it first.
     createTld("tld");
     PremiumList pl =
         persistPremiumList(

javatests/google/registry/model/server/LockTest.java

@@ -34,6 +34,7 @@ import google.registry.testing.InjectRule;
 import google.registry.util.RequestStatusChecker;
 import java.util.Optional;
 import org.joda.time.Duration;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -50,6 +51,8 @@ public class LockTest {
   private static final RequestStatusChecker requestStatusChecker = mock(RequestStatusChecker.class);
   private static final FakeClock clock = new FakeClock();
 
+  private LockMetrics origLockMetrics;
+
   @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
   @Rule public final InjectRule inject = new InjectRule();
 
@@ -73,11 +76,17 @@ public class LockTest {
 
   @Before public void setUp() {
     inject.setStaticField(Ofy.class, "clock", clock);
+    origLockMetrics = Lock.lockMetrics;
     Lock.lockMetrics = null;
     when(requestStatusChecker.getLogId()).thenReturn("current-request-id");
     when(requestStatusChecker.isRunning("current-request-id")).thenReturn(true);
   }
 
+  @After
+  public void restoreLockMetric() {
+    Lock.lockMetrics = origLockMetrics;
+  }
+
   @Test
   public void testReleasedExplicitly() {
     Optional<Lock> lock = acquire("", ONE_DAY, FREE);

javatests/google/registry/testing/TestCacheRule.java

@@ -0,0 +1,118 @@
+// 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.testing;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Maps;
+import google.registry.model.EppResource;
+import google.registry.model.index.ForeignKeyIndex;
+import google.registry.model.registry.label.PremiumList;
+import java.util.Map;
+import java.util.Optional;
+import org.joda.time.Duration;
+import org.junit.rules.ExternalResource;
+
+/**
+ * Sets up caches with desired data expiry for testing and restores their default configurations
+ * when tests complete.
+ *
+ * <p>This rule is necessary because many caches in the system are singleton and referenced through
+ * static fields.
+ */
+public class TestCacheRule extends ExternalResource {
+
+  private final ImmutableList<TestCacheHandler> cacheHandlers;
+
+  private TestCacheRule(ImmutableList<TestCacheHandler> cacheHandlers) {
+    this.cacheHandlers = cacheHandlers;
+  }
+
+  @Override
+  protected void before() {
+    cacheHandlers.forEach(TestCacheHandler::before);
+  }
+
+  @Override
+  protected void after() {
+    cacheHandlers.forEach(TestCacheHandler::after);
+  }
+
+  /** Builder for {@link TestCacheRule}. */
+  public static class Builder {
+    private final Map<String, TestCacheHandler> cacheHandlerMap = Maps.newHashMap();
+
+    public Builder withEppResourceCache(Duration expiry) {
+      cacheHandlerMap.put(
+          "EppResource.cacheEppResources",
+          new TestCacheHandler(EppResource::setCacheForTest, expiry));
+      return this;
+    }
+
+    public Builder withForeignIndexKeyCache(Duration expiry) {
+      cacheHandlerMap.put(
+          "ForeignKeyIndex.cacheForeignKeyIndexes",
+          new TestCacheHandler(ForeignKeyIndex::setCacheForTest, expiry));
+      return this;
+    }
+
+    public Builder withPremiumListsCache(Duration expiry) {
+      cacheHandlerMap.put(
+          "PremiumList.cachePremiumLists",
+          new TestCacheHandler(PremiumList::setPremiumListCacheForTest, expiry));
+      return this;
+    }
+
+    public Builder withPremiumListEntriesCache(Duration expiry) {
+      cacheHandlerMap.put(
+          "PremiumList.cachePremiumListEntries",
+          new TestCacheHandler(PremiumList::setPremiumListEntriesCacheForTest, expiry));
+      return this;
+    }
+
+    public TestCacheRule build() {
+      TestCacheRule rule = new TestCacheRule(ImmutableList.copyOf(cacheHandlerMap.values()));
+      return rule;
+    }
+  }
+
+  static class TestCacheHandler {
+    private final TestCacheSetter setter;
+    private final Duration testExpiry;
+
+    private TestCacheHandler(TestCacheSetter setter, Duration testExpiry) {
+      this.setter = setter;
+      this.testExpiry = testExpiry;
+    }
+
+    void before() {
+      setter.setCache(Optional.of(testExpiry));
+    }
+
+    void after() {
+      setter.setCache(Optional.empty());
+    }
+  }
+
+  @FunctionalInterface
+  interface TestCacheSetter {
+
+    /**
+     * Creates a new cache for use during tests.
+     *
+     * @param expiry expiry for the test cache. If not present, use default setting
+     */
+    void setCache(Optional<Duration> expiry);
+  }
+}

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

@@ -14,6 +14,7 @@
 
 package google.registry.tools.server;
 
+import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.createTlds;
 import static google.registry.testing.DatastoreHelper.persistActiveDomain;
@@ -30,6 +31,7 @@ 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.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -44,16 +46,23 @@ public class RefreshDnsForAllDomainsActionTest
   @Rule public final InjectRule inject = new InjectRule();
 
   private final DnsQueue dnsQueue = mock(DnsQueue.class);
+  private DnsQueue origDnsQueue;
 
   @Before
   public void init() {
-    inject.setStaticField(RefreshDnsForAllDomainsActionMapper.class, "dnsQueue", dnsQueue);
+    origDnsQueue = RefreshDnsForAllDomainsActionMapper.setDnsQueueForTest(dnsQueue);
 
     action = new RefreshDnsForAllDomainsAction();
     action.mrRunner = makeDefaultRunner();
     action.response = new FakeResponse();
   }
 
+  @After
+  public void restoreDnsQueue() {
+    assertThat(RefreshDnsForAllDomainsActionMapper.setDnsQueueForTest(origDnsQueue))
+        .isEqualTo(dnsQueue);
+  }
+
   private void runMapreduce() throws Exception {
     action.run();
     executeTasksUntilEmpty("mapreduce");

javatests/google/registry/util/TaskQueueUtilsTest.java

@@ -33,6 +33,7 @@ import google.registry.testing.AppEngineRule;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeSleeper;
 import org.joda.time.DateTime;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -49,6 +50,8 @@ public final class TaskQueueUtilsTest {
   public final AppEngineRule appEngine =
       AppEngineRule.builder().withDatastore().withTaskQueue().build();
 
+  private int origBatchSize;
+
   private final FakeClock clock = new FakeClock(DateTime.parse("2000-01-01TZ"));
   private final FakeSleeper sleeper = new FakeSleeper(clock);
   private final TaskQueueUtils taskQueueUtils =
@@ -59,9 +62,15 @@ public final class TaskQueueUtilsTest {
 
   @Before
   public void before() {
+    origBatchSize = TaskQueueUtils.BATCH_SIZE;
     TaskQueueUtils.BATCH_SIZE = 2;
   }
 
+  @After
+  public void after() {
+    TaskQueueUtils.BATCH_SIZE = origBatchSize;
+  }
+
   @Test
   public void testEnqueue_worksOnFirstTry_doesntSleep() {
     when(queue.add(ImmutableList.of(task))).thenReturn(ImmutableList.of(handle));

javatests/google/registry/whois/WhoisActionTest.java

@@ -31,7 +31,6 @@ import static google.registry.testing.FullFieldsTestEntityHelper.makeHostResourc
 import static google.registry.testing.FullFieldsTestEntityHelper.makeRegistrar;
 import static google.registry.testing.FullFieldsTestEntityHelper.makeRegistrarContacts;
 import static google.registry.whois.WhoisTestData.loadFile;
-import static java.util.concurrent.TimeUnit.DAYS;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_NOT_FOUND;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
@@ -43,15 +42,12 @@ import static org.mockito.Mockito.when;
 
 import com.google.appengine.api.datastore.DatastoreFailureException;
 import com.google.appengine.api.datastore.DatastoreTimeoutException;
-import com.google.common.cache.CacheBuilder;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
-import google.registry.model.EppResource;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.host.HostResource;
-import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.ofy.Ofy;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registry.Registry;
@@ -62,12 +58,14 @@ import google.registry.testing.FakeClock;
 import google.registry.testing.FakeResponse;
 import google.registry.testing.FakeSleeper;
 import google.registry.testing.InjectRule;
+import google.registry.testing.TestCacheRule;
 import google.registry.util.Retrier;
 import google.registry.whois.WhoisMetrics.WhoisMetric;
 import java.io.IOException;
 import java.io.Reader;
 import java.io.StringReader;
 import org.joda.time.DateTime;
+import org.joda.time.Duration;
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Rule;
@@ -82,6 +80,13 @@ public class WhoisActionTest {
   @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
   @Rule public final InjectRule inject = new InjectRule();
 
+  @Rule
+  public final TestCacheRule testCacheRule =
+      new TestCacheRule.Builder()
+          .withEppResourceCache(Duration.standardDays(1))
+          .withForeignIndexKeyCache(Duration.standardDays(1))
+          .build();
+
   private final FakeResponse response = new FakeResponse();
   private FakeClock clock;
 
@@ -104,10 +109,6 @@ public class WhoisActionTest {
     clock = new FakeClock(DateTime.parse("2009-06-29T20:13:00Z"));
     createTlds("lol", "xn--q9jyb4c", "1.test");
     inject.setStaticField(Ofy.class, "clock", clock);
-
-    // Set caches with long intervals, to test caching.
-    EppResource.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
-    ForeignKeyIndex.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
   }
 
   @Test