diff --git a/core/src/main/java/google/registry/tools/SetDatabaseTransitionScheduleCommand.java b/core/src/main/java/google/registry/tools/SetDatabaseTransitionScheduleCommand.java index 0885cf461..55e0e35f5 100644 --- a/core/src/main/java/google/registry/tools/SetDatabaseTransitionScheduleCommand.java +++ b/core/src/main/java/google/registry/tools/SetDatabaseTransitionScheduleCommand.java @@ -14,6 +14,8 @@ package google.registry.tools; +import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm; + import com.beust.jcommander.Parameter; import com.beust.jcommander.Parameters; import com.google.common.collect.ImmutableSortedMap; @@ -23,14 +25,13 @@ import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTr import google.registry.model.common.DatabaseTransitionSchedule.TransitionId; import google.registry.model.common.TimedTransitionProperty; import google.registry.tools.params.TransitionListParameter.PrimaryDatabaseTransitions; -import java.util.Optional; import org.joda.time.DateTime; /** Command to update {@link DatabaseTransitionSchedule}. */ @Parameters( separators = " =", commandDescription = "Set the database transition schedule for transition id.") -public class SetDatabaseTransitionScheduleCommand extends MutatingCommand { +public class SetDatabaseTransitionScheduleCommand extends ConfirmingCommand { @Parameter( names = "--transition_schedule", @@ -48,16 +49,20 @@ public class SetDatabaseTransitionScheduleCommand extends MutatingCommand { private TransitionId transitionId; @Override - protected void init() { - Optional currentSchedule = - DatabaseTransitionSchedule.get(transitionId); + protected String prompt() { + return String.format( + "Insert new schedule %s for transition ID %s?", transitionSchedule, transitionId); + } + @Override + protected String execute() { DatabaseTransitionSchedule newSchedule = DatabaseTransitionSchedule.create( transitionId, TimedTransitionProperty.fromValueMap( transitionSchedule, PrimaryDatabaseTransition.class)); - - stageEntityChange(currentSchedule.orElse(null), newSchedule); + ofyTm().transact(() -> ofyTm().put(newSchedule)); + return String.format( + "Inserted new schedule %s for transition ID %s.", transitionSchedule, transitionId); } } diff --git a/core/src/test/java/google/registry/tools/SetDatabaseTransitionScheduleCommandTest.java b/core/src/test/java/google/registry/tools/SetDatabaseTransitionScheduleCommandTest.java index b52a465ba..e4315cf87 100644 --- a/core/src/test/java/google/registry/tools/SetDatabaseTransitionScheduleCommandTest.java +++ b/core/src/test/java/google/registry/tools/SetDatabaseTransitionScheduleCommandTest.java @@ -16,7 +16,6 @@ package google.registry.tools; import static com.google.common.truth.Truth.assertThat; import static google.registry.model.common.EntityGroupRoot.getCrossTldKey; -import static google.registry.model.ofy.ObjectifyService.ofy; import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.util.DateTimeUtils.START_OF_TIME; @@ -24,12 +23,14 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import com.beust.jcommander.ParameterException; import com.google.common.collect.ImmutableSortedMap; +import com.google.common.truth.Truth8; import com.googlecode.objectify.Key; import google.registry.model.common.DatabaseTransitionSchedule; import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabase; import google.registry.model.common.DatabaseTransitionSchedule.PrimaryDatabaseTransition; import google.registry.model.common.DatabaseTransitionSchedule.TransitionId; import google.registry.model.common.TimedTransitionProperty; +import google.registry.persistence.VKey; import org.joda.time.DateTime; import org.joda.time.Duration; import org.junit.jupiter.api.BeforeEach; @@ -39,11 +40,8 @@ import org.junit.jupiter.api.Test; public class SetDatabaseTransitionScheduleCommandTest extends CommandTestCase { - Key key; - @BeforeEach void setup() { - key = Key.create(getCrossTldKey(), DatabaseTransitionSchedule.class, "test"); fakeClock.setTo(DateTime.parse("2020-12-01T00:00:00Z")); } @@ -58,7 +56,13 @@ public class SetDatabaseTransitionScheduleCommandTest @Test void testSuccess_currentScheduleIsEmpty() throws Exception { - assertThat(ofy().load().key(key).now()).isNull(); + Truth8.assertThat( + ofyTm() + .loadByKeyIfPresent( + VKey.createOfy( + DatabaseTransitionSchedule.class, + Key.create(getCrossTldKey(), DatabaseTransitionSchedule.class, "test")))) + .isEmpty(); runCommandForced( "--transition_id=SIGNED_MARK_REVOCATION_LIST", "--transition_schedule=1970-01-01T00:00:00.000Z=DATASTORE"); @@ -70,7 +74,10 @@ public class SetDatabaseTransitionScheduleCommandTest .get() .getPrimaryDatabase())) .isEqualTo(PrimaryDatabase.DATASTORE); - assertThat(command.prompt()).contains("Create DatabaseTransitionSchedule"); + assertThat(command.prompt()) + .isEqualTo( + "Insert new schedule {1970-01-01T00:00:00.000Z=DATASTORE} " + + "for transition ID SIGNED_MARK_REVOCATION_LIST?"); } @Test @@ -92,7 +99,8 @@ public class SetDatabaseTransitionScheduleCommandTest .isEqualTo(transitionMap); runCommandForced( "--transition_id=SIGNED_MARK_REVOCATION_LIST", - "--transition_schedule=1970-01-01T00:00:00.000Z=DATASTORE,2020-11-30T00:00:00.000Z=CLOUD_SQL,2020-12-06T00:00:00.000Z=DATASTORE"); + "--transition_schedule=1970-01-01T00:00:00.000Z=DATASTORE," + + "2020-11-30T00:00:00.000Z=CLOUD_SQL,2020-12-06T00:00:00.000Z=DATASTORE"); ImmutableSortedMap retrievedTransitionMap = ofyTm() .transact( @@ -117,6 +125,10 @@ public class SetDatabaseTransitionScheduleCommandTest .get() .getPrimaryDatabase())) .isEqualTo(PrimaryDatabase.DATASTORE); - assertThat(command.prompt()).contains("Update DatabaseTransitionSchedule"); + assertThat(command.prompt()) + .isEqualTo( + "Insert new schedule {1970-01-01T00:00:00.000Z=DATASTORE, " + + "2020-11-30T00:00:00.000Z=CLOUD_SQL, 2020-12-06T00:00:00.000Z=DATASTORE} " + + "for transition ID SIGNED_MARK_REVOCATION_LIST?"); } }