Disable JpaEntityCoverageCheck by default (#555)

* Disable JpaEntityCoverageCheck by default

Only members of SqlIntegrationTestSuite should enable the check,
which incurs per-test overhead.

core/src/main/java/google/registry/tools/CheckDomainCommand.java

@@ -14,6 +14,8 @@
 
 package google.registry.tools;
 
+import static com.google.common.base.Strings.isNullOrEmpty;
+
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.google.common.collect.Multimap;
@@ -39,6 +41,11 @@ final class CheckDomainCommand extends NonMutatingEppToolCommand {
       required = true)
   private List<String> mainParameters;
 
+  @Parameter(
+      names = {"-t", "--token"},
+      description = "Allocation token to use in the command, if desired")
+  private String allocationToken;
+
   @Inject
   @Config("registryAdminClientId")
   String registryAdminClientId;
@@ -53,7 +60,11 @@ final class CheckDomainCommand extends NonMutatingEppToolCommand {
     Multimap<String, String> domainNameMap = validateAndGroupDomainNamesByTld(mainParameters);
     for (Collection<String> values : domainNameMap.asMap().values()) {
       setSoyTemplate(DomainCheckSoyInfo.getInstance(), DomainCheckSoyInfo.DOMAINCHECK);
-      addSoyRecord(clientId, new SoyMapData("domainNames", values));
+      SoyMapData soyMapData = new SoyMapData("domainNames", values);
+      if (!isNullOrEmpty(allocationToken)) {
+        soyMapData.put("allocationToken", allocationToken);
+      }
+      addSoyRecord(clientId, soyMapData);
     }
   }
 }

core/src/main/resources/google/registry/tools/soy/DomainCheck.soy

@@ -19,6 +19,7 @@
  */
 {template .domaincheck stricthtml="false"}
 {@param domainNames: list<string>}
+{@param? allocationToken: string|null}
 <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
 <epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
   <command>
@@ -39,6 +40,12 @@
         </fee:domain>
       {/for}
       </fee:check>
+      {if isNonnull($allocationToken)}
+        <allocationToken:allocationToken
+            xmlns:allocationToken="urn:ietf:params:xml:ns:allocationToken-1.0">
+          {$allocationToken}
+        </allocationToken:allocationToken>
+      {/if}
     </extension>
     <clTRID>RegistryTool</clTRID>
   </command>

core/src/test/java/google/registry/model/EntityTestCase.java

@@ -50,12 +50,23 @@ public abstract class EntityTestCase {
 
   protected FakeClock fakeClock = new FakeClock(DateTime.now(UTC));
 
-  @Rule @RegisterExtension
-  public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+  @Rule @RegisterExtension public final AppEngineRule appEngine;
 
   @Rule @RegisterExtension public InjectRule inject = new InjectRule();
 
+  protected EntityTestCase() {
+    this(false);
+  }
+
+  protected EntityTestCase(boolean enableJpaEntityCheck) {
+    appEngine =
+        AppEngineRule.builder()
+            .withDatastoreAndCloudSql()
+            .enableJpaEntityCoverageCheck(enableJpaEntityCheck)
+            .withClock(fakeClock)
+            .build();
+  }
+
   @Before
   @BeforeEach
   public void injectClock() {

core/src/test/java/google/registry/model/domain/DomainBaseSqlTest.java

@@ -38,6 +38,10 @@ public class DomainBaseSqlTest extends EntityTestCase {
   Key<ContactResource> contactKey;
   Key<ContactResource> contact2Key;
 
+  public DomainBaseSqlTest() {
+    super(true);
+  }
+
   @BeforeEach
   public void setUp() {
     contactKey = Key.create(ContactResource.class, "contact_id1");

core/src/test/java/google/registry/model/registry/RegistryLockDaoTest.java

@@ -41,7 +41,11 @@ public final class RegistryLockDaoTest {
 
   @RegisterExtension
   public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .enableJpaEntityCoverageCheck(true)
+          .withClock(fakeClock)
+          .build();
 
   @Test
   public void testSaveAndLoad_success() {

core/src/test/java/google/registry/schema/cursor/CursorDaoTest.java

@@ -43,7 +43,11 @@ public class CursorDaoTest {
 
   @RegisterExtension
   public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .enableJpaEntityCoverageCheck(true)
+          .withClock(fakeClock)
+          .build();
 
   @Test
   public void save_worksSuccessfullyOnNewCursor() {

core/src/test/java/google/registry/schema/registrar/RegistrarDaoTest.java

@@ -33,6 +33,10 @@ public class RegistrarDaoTest extends EntityTestCase {
 
   private Registrar testRegistrar;
 
+  public RegistrarDaoTest() {
+    super(true);
+  }
+
   @BeforeEach
   public void setUp() {
     testRegistrar =

core/src/test/java/google/registry/schema/tld/PremiumListDaoTest.java

@@ -45,7 +45,11 @@ public class PremiumListDaoTest {
 
   @RegisterExtension
   public final AppEngineRule appEngine =
-      AppEngineRule.builder().withDatastoreAndCloudSql().withClock(fakeClock).build();
+      AppEngineRule.builder()
+          .withDatastoreAndCloudSql()
+          .enableJpaEntityCoverageCheck(true)
+          .withClock(fakeClock)
+          .build();
 
   private static final ImmutableMap<String, BigDecimal> TEST_PRICES =
       ImmutableMap.of(

core/src/test/java/google/registry/testing/AppEngineRule.java

@@ -41,6 +41,7 @@ import google.registry.model.registrar.Registrar.State;
 import google.registry.model.registrar.RegistrarAddress;
 import google.registry.model.registrar.RegistrarContact;
 import google.registry.persistence.transaction.JpaTestRules;
+import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationTestRule;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageExtension;
 import google.registry.persistence.transaction.JpaTestRules.JpaIntegrationWithCoverageRule;
 import google.registry.util.Clock;
@@ -104,10 +105,19 @@ public final class AppEngineRule extends ExternalResource
 
   /** A rule-within-a-rule to provide a temporary folder for AppEngineRule's internal temp files. */
   TemporaryFolder temporaryFolder = new TemporaryFolder();
-  // Sets up a SQL database when running on JUnit 5.
+  /**
+   * Sets up a SQL database when running on JUnit 5. This is for test classes that are not member of
+   * the {@code SqlIntegrationTestSuite}.
+   */
+  JpaIntegrationTestRule jpaIntegrationTestRule = null;
+  /**
+   * Sets up a SQL database when running on JUnit 5 and records the JPA entities tested by each test
+   * class. This is for {@code SqlIntegrationTestSuite} members.
+   */
   JpaIntegrationWithCoverageExtension jpaIntegrationWithCoverageExtension = null;
 
   private boolean withDatastoreAndCloudSql;
+  private boolean enableJpaEntityCoverageCheck;
   private boolean withLocalModules;
   private boolean withTaskQueue;
   private boolean withUserService;
@@ -127,6 +137,14 @@ public final class AppEngineRule extends ExternalResource
       rule.withDatastoreAndCloudSql = true;
       return this;
     }
+    /**
+     * Enables JPA entity coverage check if {@code enabled} is true. This should only be enabled for
+     * members of SqlIntegrationTestSuite.
+     */
+    public Builder enableJpaEntityCoverageCheck(boolean enabled) {
+      rule.enableJpaEntityCoverageCheck = enabled;
+      return this;
+    }
 
     /** Turn on the use of local modules. */
     public Builder withLocalModules() {
@@ -164,6 +182,9 @@ public final class AppEngineRule extends ExternalResource
     }
 
     public AppEngineRule build() {
+      checkState(
+          !rule.enableJpaEntityCoverageCheck || rule.withDatastoreAndCloudSql,
+          "withJpaEntityCoverageCheck enabled without Cloud SQL");
       return rule;
     }
   }
@@ -279,8 +300,13 @@ public final class AppEngineRule extends ExternalResource
       if (clock != null) {
         builder.withClock(clock);
       }
-      jpaIntegrationWithCoverageExtension = builder.buildIntegrationWithCoverageExtension();
-      jpaIntegrationWithCoverageExtension.beforeEach(context);
+      if (enableJpaEntityCoverageCheck) {
+        jpaIntegrationWithCoverageExtension = builder.buildIntegrationWithCoverageExtension();
+        jpaIntegrationWithCoverageExtension.beforeEach(context);
+      } else {
+        jpaIntegrationTestRule = builder.buildIntegrationTestRule();
+        jpaIntegrationTestRule.before();
+      }
     }
   }
 
@@ -288,9 +314,11 @@ public final class AppEngineRule extends ExternalResource
   @Override
   public void afterEach(ExtensionContext context) throws Exception {
     if (withDatastoreAndCloudSql) {
-      checkState(
-          jpaIntegrationWithCoverageExtension != null, "Null jpaIntegrationWithCoverageExtension");
-      jpaIntegrationWithCoverageExtension.afterEach(context);
+      if (enableJpaEntityCoverageCheck) {
+        jpaIntegrationWithCoverageExtension.afterEach(context);
+      } else {
+        jpaIntegrationTestRule.after();
+      }
     }
     after();
   }
@@ -310,7 +338,10 @@ public final class AppEngineRule extends ExternalResource
       if (clock != null) {
         builder.withClock(clock);
       }
-      statement = builder.buildIntegrationWithCoverageRule().apply(base, description);
+      checkState(
+          !enableJpaEntityCoverageCheck,
+          "JUnit4 tests must not enable withJpaEntityCoverageCheck.");
+      statement = builder.buildIntegrationTestRule().apply(base, description);
     }
     return super.apply(statement, description);
   }

core/src/test/java/google/registry/tools/CheckDomainCommandTest.java

@@ -76,6 +76,12 @@ public class CheckDomainCommandTest extends EppToolCommandTestCase<CheckDomainCo
     eppVerifier.expectDryRun().expectClientId("adminreg").verifySent("domain_check_fee.xml");
   }
 
+  @Test
+  public void testSuccess_allocationToken_reserved() throws Exception {
+    runCommand("--client=NewRegistrar", "--token=abc123", "example.tld");
+    eppVerifier.expectDryRun().verifySent("domain_check_fee_allocationtoken.xml");
+  }
+
   @Test
   public void testFailure_NoMainParameter() {
     assertThrows(ParameterException.class, () -> runCommand("--client=NewRegistrar"));

core/src/test/resources/google/registry/tools/server/domain_check_fee_allocationtoken.xml

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
+  <command>
+    <check>
+      <domain:check xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
+        <domain:name>example.tld</domain:name>
+      </domain:check>
+    </check>
+    <extension>
+      <fee:check xmlns:fee="urn:ietf:params:xml:ns:fee-0.6">
+        <fee:domain>
+          <fee:name>example.tld</fee:name>
+          <fee:command>create</fee:command>
+          <fee:period unit="y">1</fee:period>
+        </fee:domain>
+      </fee:check>
+      <allocationToken:allocationToken
+          xmlns:allocationToken="urn:ietf:params:xml:ns:allocationToken-1.0">
+        abc123
+      </allocationToken:allocationToken>
+    </extension>
+    <clTRID>RegistryTool</clTRID>
+  </command>
+</epp>