From 4df4bf148944d25de62c7ad6aef2909d9ec40ade Mon Sep 17 00:00:00 2001 From: Ben McIlwain Date: Tue, 16 Jun 2026 17:24:01 -0400 Subject: [PATCH] Add FeatureFlag helper methods to DatabaseHelper (#3087) Introduced overloaded helper methods `persistFeatureFlag` in `DatabaseHelper` to simplify the creation and persistence of `FeatureFlag` entities in test setups. 1. `persistFeatureFlag(FeatureName, FeatureStatus)`: Persists a feature flag with a single status starting from the Unix Epoch (`START_INSTANT`). 2. `persistFeatureFlag(FeatureName, FeatureStatus, Instant, FeatureStatus)`: Persists a feature flag with an initial status at `START_INSTANT` and a subsequent transition at a specified time. Refactored 11 occurrences of manual 1-transition flag creation and 5 occurrences of 2-transition flag creation across the test suite to use these new helpers, significantly reducing boilerplate and improving test readability. TAG=agy CONV=583b8a23-9fe5-476d-ac35-aeba7b218eb0 --- .../export/ExportDomainListsActionTest.java | 23 +++------ .../flows/domain/DomainCreateFlowTest.java | 26 ++-------- .../flows/domain/DomainUpdateFlowTest.java | 32 +++--------- .../flows/session/LoginFlowTestCase.java | 10 +--- .../registry/testing/DatabaseHelper.java | 29 +++++++++++ .../ConfigureFeatureFlagCommandTest.java | 30 ++---------- .../tools/CreateDomainCommandTest.java | 10 +--- .../tools/DeleteFeatureFlagCommandTest.java | 11 +---- .../tools/GetFeatureFlagCommandTest.java | 48 ++++-------------- .../tools/ListFeatureFlagsCommandTest.java | 49 +++++-------------- .../tools/UpdateDomainCommandTest.java | 10 +--- 11 files changed, 79 insertions(+), 199 deletions(-) diff --git a/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java b/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java index c4d1f0408..193aeae72 100644 --- a/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java +++ b/core/src/test/java/google/registry/export/ExportDomainListsActionTest.java @@ -21,8 +21,8 @@ import static google.registry.model.common.FeatureFlag.FeatureStatus.INACTIVE; import static google.registry.testing.DatabaseHelper.createTld; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistDeletedDomain; +import static google.registry.testing.DatabaseHelper.persistFeatureFlag; import static google.registry.testing.DatabaseHelper.persistResource; -import static google.registry.util.DateTimeUtils.START_INSTANT; import static google.registry.util.DateTimeUtils.plusDays; import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -35,10 +35,8 @@ import com.google.cloud.storage.BlobId; import com.google.cloud.storage.StorageException; import com.google.cloud.storage.contrib.nio.testing.LocalStorageHelper; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableSortedMap; import com.google.common.net.MediaType; import google.registry.gcs.GcsUtils; -import google.registry.model.common.FeatureFlag; import google.registry.model.domain.Domain; import google.registry.model.domain.GracePeriod; import google.registry.model.domain.rgp.GracePeriodStatus; @@ -80,7 +78,7 @@ class ExportDomainListsActionTest { action.gcsUtils = gcsUtils; action.clock = clock; action.driveConnection = driveConnection; - persistFeatureFlag(INACTIVE); + persistFeatureFlag(INCLUDE_PENDING_DELETE_DATE_FOR_DOMAINS, INACTIVE); } private void verifyExportedToDrive(String folderId, String filename, String domains) @@ -110,7 +108,7 @@ class ExportDomainListsActionTest { @Test void test_outputsOnlyActiveDomains_csv() throws Exception { - persistFeatureFlag(ACTIVE); + persistFeatureFlag(INCLUDE_PENDING_DELETE_DATE_FOR_DOMAINS, ACTIVE); persistActiveDomain("onetwo.tld"); persistActiveDomain("rudnitzky.tld"); persistDeletedDomain("mortuary.tld", Instant.parse("2001-03-14T10:11:12Z")); @@ -144,7 +142,7 @@ class ExportDomainListsActionTest { @Test void test_outputsOnlyDomainsOnRealTlds_csv() throws Exception { - persistFeatureFlag(ACTIVE); + persistFeatureFlag(INCLUDE_PENDING_DELETE_DATE_FOR_DOMAINS, ACTIVE); persistActiveDomain("onetwo.tld"); persistActiveDomain("rudnitzky.tld"); persistActiveDomain("wontgo.testtld"); @@ -164,7 +162,7 @@ class ExportDomainListsActionTest { @Test void test_outputIncludesDeletionTimes_forPendingDeletes_notRdemption() throws Exception { - persistFeatureFlag(ACTIVE); + persistFeatureFlag(INCLUDE_PENDING_DELETE_DATE_FOR_DOMAINS, ACTIVE); // Domains pending delete (meaning the 5 day period, not counting the 30 day redemption period) // should include their pending deletion date persistActiveDomain("active.tld"); @@ -227,7 +225,7 @@ class ExportDomainListsActionTest { @Test void test_outputsDomainsFromDifferentTldsToMultipleFiles_csv() throws Exception { - persistFeatureFlag(ACTIVE); + persistFeatureFlag(INCLUDE_PENDING_DELETE_DATE_FOR_DOMAINS, ACTIVE); createTld("tldtwo"); persistResource(Tld.get("tldtwo").asBuilder().setDriveFolderId("hooray").build()); @@ -255,13 +253,4 @@ class ExportDomainListsActionTest { // tldthree does not have a drive id, so no export to drive is performed. verifyNoMoreInteractions(driveConnection); } - - private void persistFeatureFlag(FeatureFlag.FeatureStatus status) { - persistResource( - new FeatureFlag() - .asBuilder() - .setFeatureName(INCLUDE_PENDING_DELETE_DATE_FOR_DOMAINS) - .setStatusMap(ImmutableSortedMap.of(START_INSTANT, status)) - .build()); - } } diff --git a/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java b/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java index c86be82a1..2265cf8c4 100644 --- a/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java +++ b/core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java @@ -53,6 +53,7 @@ import static google.registry.testing.DatabaseHelper.loadRegistrar; import static google.registry.testing.DatabaseHelper.newHost; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistActiveHost; +import static google.registry.testing.DatabaseHelper.persistFeatureFlag; import static google.registry.testing.DatabaseHelper.persistReservedList; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.testing.DomainSubject.assertAboutDomains; @@ -151,7 +152,6 @@ import google.registry.model.billing.BillingBase.Reason; import google.registry.model.billing.BillingBase.RenewalPriceBehavior; import google.registry.model.billing.BillingEvent; import google.registry.model.billing.BillingRecurrence; -import google.registry.model.common.FeatureFlag; import google.registry.model.common.FeatureFlag.FeatureStatus; import google.registry.model.domain.Domain; import google.registry.model.domain.DomainHistory; @@ -968,11 +968,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase { sessionMetadata.setRegistrarId(null); // Don't implicitly log in (all other flows need to). registrar = loadRegistrar("NewRegistrar"); registrarBuilder = registrar.asBuilder(); - persistResource( - new FeatureFlag.Builder() - .setFeatureName(PROHIBIT_CONTACT_OBJECTS_ON_LOGIN) - .setStatusMap(ImmutableSortedMap.of(START_INSTANT, FeatureStatus.ACTIVE)) - .build()); + persistFeatureFlag(PROHIBIT_CONTACT_OBJECTS_ON_LOGIN, FeatureStatus.ACTIVE); } // Can't inline this since it may be overridden in subclasses. diff --git a/core/src/test/java/google/registry/testing/DatabaseHelper.java b/core/src/test/java/google/registry/testing/DatabaseHelper.java index 024eb7d59..d82688456 100644 --- a/core/src/test/java/google/registry/testing/DatabaseHelper.java +++ b/core/src/test/java/google/registry/testing/DatabaseHelper.java @@ -69,6 +69,9 @@ import google.registry.model.billing.BillingCancellation; import google.registry.model.billing.BillingEvent; import google.registry.model.billing.BillingRecurrence; import google.registry.model.common.DnsRefreshRequest; +import google.registry.model.common.FeatureFlag; +import google.registry.model.common.FeatureFlag.FeatureName; +import google.registry.model.common.FeatureFlag.FeatureStatus; import google.registry.model.console.GlobalRole; import google.registry.model.console.User; import google.registry.model.console.UserRoles; @@ -682,6 +685,32 @@ public final class DatabaseHelper { return newRegistrars.build(); } + /** Persists and returns a {@link FeatureFlag} with a single status starting from the epoch. */ + public static FeatureFlag persistFeatureFlag(FeatureName featureName, FeatureStatus status) { + return persistFeatureFlag(featureName, ImmutableSortedMap.of(START_INSTANT, status)); + } + + /** Persists and returns a {@link FeatureFlag} with an initial status and one transition. */ + public static FeatureFlag persistFeatureFlag( + FeatureName featureName, + FeatureStatus initialStatus, + Instant transitionTime, + FeatureStatus status) { + return persistFeatureFlag( + featureName, + ImmutableSortedMap.naturalOrder() + .put(START_INSTANT, initialStatus) + .put(transitionTime, status) + .build()); + } + + /** Persists and returns a {@link FeatureFlag} with a custom status map. */ + public static FeatureFlag persistFeatureFlag( + FeatureName featureName, ImmutableSortedMap statusMap) { + return persistResource( + new FeatureFlag.Builder().setFeatureName(featureName).setStatusMap(statusMap).build()); + } + public static Iterable getBillingEvents() { return tm().transact( () -> diff --git a/core/src/test/java/google/registry/tools/ConfigureFeatureFlagCommandTest.java b/core/src/test/java/google/registry/tools/ConfigureFeatureFlagCommandTest.java index 212b08af8..e90a7e958 100644 --- a/core/src/test/java/google/registry/tools/ConfigureFeatureFlagCommandTest.java +++ b/core/src/test/java/google/registry/tools/ConfigureFeatureFlagCommandTest.java @@ -19,7 +19,7 @@ import static google.registry.model.common.FeatureFlag.FeatureName.MINIMUM_DATAS import static google.registry.model.common.FeatureFlag.FeatureName.TEST_FEATURE; import static google.registry.model.common.FeatureFlag.FeatureStatus.ACTIVE; import static google.registry.model.common.FeatureFlag.FeatureStatus.INACTIVE; -import static google.registry.testing.DatabaseHelper.persistResource; +import static google.registry.testing.DatabaseHelper.persistFeatureFlag; import static google.registry.util.DateTimeUtils.START_INSTANT; import static google.registry.util.DateTimeUtils.plusWeeks; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -27,7 +27,6 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import com.beust.jcommander.ParameterException; import com.google.common.collect.ImmutableSortedMap; import google.registry.model.common.FeatureFlag; -import google.registry.model.common.FeatureFlag.FeatureStatus; import google.registry.model.common.TimedTransitionProperty; import google.registry.testing.FakeClock; import java.time.Instant; @@ -81,14 +80,7 @@ public class ConfigureFeatureFlagCommandTest extends CommandTestCasenaturalOrder() - .put(START_INSTANT, INACTIVE) - .build()) - .build()); + persistFeatureFlag(TEST_FEATURE, INACTIVE); Instant featureStart = plusWeeks(clock.now(), 6); assertThat(FeatureFlag.get(TEST_FEATURE).getStatusMap()) @@ -110,14 +102,7 @@ public class ConfigureFeatureFlagCommandTest extends CommandTestCasenaturalOrder() - .put(START_INSTANT, INACTIVE) - .build()) - .build()); + persistFeatureFlag(TEST_FEATURE, INACTIVE); Instant featureStart = plusWeeks(clock.now(), 6); assertThat(FeatureFlag.get(TEST_FEATURE).getStatusMap()) @@ -179,14 +164,7 @@ public class ConfigureFeatureFlagCommandTest extends CommandTestCasenaturalOrder() - .put(START_INSTANT, INACTIVE) - .build()) - .build()); + persistFeatureFlag(TEST_FEATURE, INACTIVE); Instant featureStart = plusWeeks(clock.now(), 6); assertThat(FeatureFlag.get(TEST_FEATURE).getStatusMap()) diff --git a/core/src/test/java/google/registry/tools/CreateDomainCommandTest.java b/core/src/test/java/google/registry/tools/CreateDomainCommandTest.java index cddb50879..7d0e3dd57 100644 --- a/core/src/test/java/google/registry/tools/CreateDomainCommandTest.java +++ b/core/src/test/java/google/registry/tools/CreateDomainCommandTest.java @@ -18,17 +18,15 @@ import static com.google.common.truth.Truth.assertThat; import static google.registry.model.common.FeatureFlag.FeatureName.FORBID_INSECURE_ALGORITHMS_RFC_9904; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.testing.DatabaseHelper.createTld; +import static google.registry.testing.DatabaseHelper.persistFeatureFlag; import static google.registry.testing.DatabaseHelper.persistPremiumList; import static google.registry.testing.DatabaseHelper.persistResource; -import static google.registry.util.DateTimeUtils.START_INSTANT; import static org.joda.money.CurrencyUnit.JPY; import static org.junit.jupiter.api.Assertions.assertThrows; import com.beust.jcommander.ParameterException; import com.google.common.collect.ImmutableSet; -import com.google.common.collect.ImmutableSortedMap; import google.registry.dns.writer.VoidDnsWriter; -import google.registry.model.common.FeatureFlag; import google.registry.model.common.FeatureFlag.FeatureStatus; import google.registry.model.pricing.StaticPremiumListPricingEngine; import google.registry.model.tld.Tld; @@ -286,11 +284,7 @@ class CreateDomainCommandTest extends EppToolCommandTestCase FeatureFlag.isActiveNow(TEST_FEATURE))).isTrue(); runCommandForced("TEST_FEATURE"); assertThat(FeatureFlag.getUncached(TEST_FEATURE)).isEmpty(); diff --git a/core/src/test/java/google/registry/tools/GetFeatureFlagCommandTest.java b/core/src/test/java/google/registry/tools/GetFeatureFlagCommandTest.java index b43dd2a2b..1e0111276 100644 --- a/core/src/test/java/google/registry/tools/GetFeatureFlagCommandTest.java +++ b/core/src/test/java/google/registry/tools/GetFeatureFlagCommandTest.java @@ -19,7 +19,7 @@ import static google.registry.model.common.FeatureFlag.FeatureName.MINIMUM_DATAS import static google.registry.model.common.FeatureFlag.FeatureName.TEST_FEATURE; import static google.registry.model.common.FeatureFlag.FeatureStatus.ACTIVE; import static google.registry.model.common.FeatureFlag.FeatureStatus.INACTIVE; -import static google.registry.testing.DatabaseHelper.persistResource; +import static google.registry.testing.DatabaseHelper.persistFeatureFlag; import static google.registry.util.DateTimeUtils.START_INSTANT; import static google.registry.util.DateTimeUtils.plusWeeks; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -27,7 +27,6 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import com.beust.jcommander.ParameterException; import com.google.common.collect.ImmutableSortedMap; import google.registry.model.EntityYamlUtils; -import google.registry.model.common.FeatureFlag; import google.registry.model.common.FeatureFlag.FeatureFlagNotFoundException; import google.registry.model.common.FeatureFlag.FeatureStatus; import google.registry.testing.FakeClock; @@ -47,15 +46,7 @@ public class GetFeatureFlagCommandTest extends CommandTestCasenaturalOrder() - .put(START_INSTANT, INACTIVE) - .put(plusWeeks(clock.now(), 8), ACTIVE) - .build()) - .build()); + persistFeatureFlag(TEST_FEATURE, INACTIVE, plusWeeks(clock.now(), 8), ACTIVE); runCommand("TEST_FEATURE"); assertInStdout( """ @@ -68,24 +59,13 @@ public class GetFeatureFlagCommandTest extends CommandTestCasenaturalOrder() - .put(START_INSTANT, INACTIVE) - .put(plusWeeks(clock.now(), 8), ACTIVE) - .build()) - .build()); - persistResource( - new FeatureFlag.Builder() - .setFeatureName(MINIMUM_DATASET_CONTACTS_OPTIONAL) - .setStatusMap( - ImmutableSortedMap.naturalOrder() - .put(START_INSTANT, INACTIVE) - .put(plusWeeks(clock.now(), 3), ACTIVE) - .put(plusWeeks(clock.now(), 6), INACTIVE) - .build()) + persistFeatureFlag(TEST_FEATURE, INACTIVE, plusWeeks(clock.now(), 8), ACTIVE); + persistFeatureFlag( + MINIMUM_DATASET_CONTACTS_OPTIONAL, + ImmutableSortedMap.naturalOrder() + .put(START_INSTANT, INACTIVE) + .put(plusWeeks(clock.now(), 3), ACTIVE) + .put(plusWeeks(clock.now(), 6), INACTIVE) .build()); runCommand("TEST_FEATURE", "MINIMUM_DATASET_CONTACTS_OPTIONAL"); assertInStdout( @@ -114,15 +94,7 @@ public class GetFeatureFlagCommandTest extends CommandTestCasenaturalOrder() - .put(START_INSTANT, INACTIVE) - .put(plusWeeks(clock.now(), 8), ACTIVE) - .build()) - .build()); + persistFeatureFlag(TEST_FEATURE, INACTIVE, plusWeeks(clock.now(), 8), ACTIVE); assertThrows( FeatureFlagNotFoundException.class, () -> runCommand("TEST_FEATURE", "MINIMUM_DATASET_CONTACTS_OPTIONAL")); diff --git a/core/src/test/java/google/registry/tools/ListFeatureFlagsCommandTest.java b/core/src/test/java/google/registry/tools/ListFeatureFlagsCommandTest.java index be1cf0102..8905079f8 100644 --- a/core/src/test/java/google/registry/tools/ListFeatureFlagsCommandTest.java +++ b/core/src/test/java/google/registry/tools/ListFeatureFlagsCommandTest.java @@ -18,14 +18,13 @@ import static google.registry.model.common.FeatureFlag.FeatureName.MINIMUM_DATAS import static google.registry.model.common.FeatureFlag.FeatureName.TEST_FEATURE; import static google.registry.model.common.FeatureFlag.FeatureStatus.ACTIVE; import static google.registry.model.common.FeatureFlag.FeatureStatus.INACTIVE; -import static google.registry.testing.DatabaseHelper.persistResource; +import static google.registry.testing.DatabaseHelper.persistFeatureFlag; import static google.registry.testing.TestDataHelper.loadFile; import static google.registry.util.DateTimeUtils.START_INSTANT; import static google.registry.util.DateTimeUtils.plusWeeks; import com.google.common.collect.ImmutableSortedMap; import google.registry.model.EntityYamlUtils; -import google.registry.model.common.FeatureFlag; import google.registry.model.common.FeatureFlag.FeatureStatus; import java.time.Instant; import org.junit.jupiter.api.BeforeEach; @@ -41,49 +40,23 @@ public class ListFeatureFlagsCommandTest extends CommandTestCasenaturalOrder() - .put(START_INSTANT, INACTIVE) - .put(plusWeeks(fakeClock.now(), 8), ACTIVE) - .build()) - .build()); + persistFeatureFlag(TEST_FEATURE, INACTIVE, plusWeeks(fakeClock.now(), 8), ACTIVE); runCommand(); assertInStdout(loadFile(getClass(), "oneFlag.yaml")); } @Test void test_success_manyFlags() throws Exception { - persistResource( - new FeatureFlag.Builder() - .setFeatureName(TEST_FEATURE) - .setStatusMap( - ImmutableSortedMap.naturalOrder() - .put(START_INSTANT, INACTIVE) - .put(plusWeeks(fakeClock.now(), 8), ACTIVE) - .build()) - .build()); - persistResource( - new FeatureFlag.Builder() - .setFeatureName(MINIMUM_DATASET_CONTACTS_OPTIONAL) - .setStatusMap( - ImmutableSortedMap.naturalOrder() - .put(START_INSTANT, INACTIVE) - .put(plusWeeks(fakeClock.now(), 1), ACTIVE) - .put(plusWeeks(fakeClock.now(), 8), INACTIVE) - .put(plusWeeks(fakeClock.now(), 10), ACTIVE) - .build()) - .build()); - persistResource( - new FeatureFlag.Builder() - .setFeatureName(MINIMUM_DATASET_CONTACTS_PROHIBITED) - .setStatusMap( - ImmutableSortedMap.naturalOrder() - .put(START_INSTANT, ACTIVE) - .build()) + persistFeatureFlag(TEST_FEATURE, INACTIVE, plusWeeks(fakeClock.now(), 8), ACTIVE); + persistFeatureFlag( + MINIMUM_DATASET_CONTACTS_OPTIONAL, + ImmutableSortedMap.naturalOrder() + .put(START_INSTANT, INACTIVE) + .put(plusWeeks(fakeClock.now(), 1), ACTIVE) + .put(plusWeeks(fakeClock.now(), 8), INACTIVE) + .put(plusWeeks(fakeClock.now(), 10), ACTIVE) .build()); + persistFeatureFlag(MINIMUM_DATASET_CONTACTS_PROHIBITED, ACTIVE); runCommand(); assertInStdout(loadFile(getClass(), "threeFlags.yaml")); } diff --git a/core/src/test/java/google/registry/tools/UpdateDomainCommandTest.java b/core/src/test/java/google/registry/tools/UpdateDomainCommandTest.java index 601929f12..d80e47698 100644 --- a/core/src/test/java/google/registry/tools/UpdateDomainCommandTest.java +++ b/core/src/test/java/google/registry/tools/UpdateDomainCommandTest.java @@ -23,11 +23,11 @@ import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE; import static google.registry.testing.DatabaseHelper.createTld; import static google.registry.testing.DatabaseHelper.persistActiveDomain; import static google.registry.testing.DatabaseHelper.persistActiveHost; +import static google.registry.testing.DatabaseHelper.persistFeatureFlag; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.testing.TestLogHandlerUtils.assertLogMessage; import static google.registry.testing.TestLogHandlerUtils.assertNoLogMessage; import static google.registry.util.DateTimeUtils.END_INSTANT; -import static google.registry.util.DateTimeUtils.START_INSTANT; import static google.registry.util.DateTimeUtils.minusDays; import static google.registry.util.DateTimeUtils.plusDays; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -35,11 +35,9 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import com.beust.jcommander.ParameterException; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; -import com.google.common.collect.ImmutableSortedMap; import google.registry.model.billing.BillingBase.Flag; import google.registry.model.billing.BillingBase.Reason; import google.registry.model.billing.BillingRecurrence; -import google.registry.model.common.FeatureFlag; import google.registry.model.common.FeatureFlag.FeatureStatus; import google.registry.model.domain.Domain; import google.registry.model.domain.DomainHistory; @@ -542,11 +540,7 @@ class UpdateDomainCommandTest extends EppToolCommandTestCase