From 43a534f05bd8d15f0795d2db3714ce53ba1994f1 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Mon, 12 May 2025 16:55:55 -0400 Subject: [PATCH] Add a command for regenerating account constraint tables --- .../textsecuregcm/WhisperServerService.java | 3 + .../textsecuregcm/storage/Accounts.java | 70 ++++++- .../workers/CommandDependencies.java | 2 + ...egenerateAccountConstraintDataCommand.java | 84 ++++++++ .../textsecuregcm/storage/AccountsTest.java | 191 +++++++++++++++++- ...PushNotificationExperimentCommandTest.java | 3 +- .../workers/NotifyIdleDevicesCommandTest.java | 1 + ...erateAccountConstraintDataCommandTest.java | 96 +++++++++ ...PushNotificationExperimentCommandTest.java | 1 + 9 files changed, 446 insertions(+), 5 deletions(-) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommand.java create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommandTest.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index ab7334da1..2f52c7dcb 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -269,6 +269,7 @@ import org.whispersystems.textsecuregcm.workers.IdleDeviceNotificationSchedulerF import org.whispersystems.textsecuregcm.workers.MessagePersisterServiceCommand; import org.whispersystems.textsecuregcm.workers.NotifyIdleDevicesCommand; import org.whispersystems.textsecuregcm.workers.ProcessScheduledJobsServiceCommand; +import org.whispersystems.textsecuregcm.workers.RegenerateAccountConstraintDataCommand; import org.whispersystems.textsecuregcm.workers.RemoveExpiredAccountsCommand; import org.whispersystems.textsecuregcm.workers.RemoveExpiredBackupsCommand; import org.whispersystems.textsecuregcm.workers.RemoveExpiredLinkedDevicesCommand; @@ -335,6 +336,8 @@ public class WhisperServerService extends Application ACCOUNT_FIELDS_TO_EXCLUDE_FROM_SERIALIZATION = List.of("uuid", "usernameLinkHandle"); - private static final ObjectWriter ACCOUNT_DDB_JSON_WRITER = SystemMapper.jsonMapper() + @VisibleForTesting + static final ObjectWriter ACCOUNT_DDB_JSON_WRITER = SystemMapper.jsonMapper() .writer(SystemMapper.excludingField(Account.class, ACCOUNT_FIELDS_TO_EXCLUDE_FROM_SERIALIZATION)); private static final Timer CREATE_TIMER = Metrics.timer(name(Accounts.class, "create")); @@ -1404,8 +1407,7 @@ public class Accounts { final String tableName, final AttributeValue uuidAttr, final String keyName, - final AttributeValue keyValue - ) { + final AttributeValue keyValue) { return TransactWriteItem.builder() .put(Put.builder() .tableName(tableName) @@ -1470,6 +1472,68 @@ public class Accounts { .build(); } + public CompletableFuture regenerateConstraints(final Account account) { + final List> constraintFutures = new ArrayList<>(); + + constraintFutures.add(writeConstraint(phoneNumberConstraintTableName, + account.getIdentifier(IdentityType.ACI), + ATTR_ACCOUNT_E164, + AttributeValues.fromString(account.getNumber()))); + + constraintFutures.add(writeConstraint(phoneNumberIdentifierConstraintTableName, + account.getIdentifier(IdentityType.ACI), + ATTR_PNI_UUID, + AttributeValues.fromUUID(account.getPhoneNumberIdentifier()))); + + account.getUsernameHash().ifPresent(usernameHash -> + constraintFutures.add(writeUsernameConstraint(account.getIdentifier(IdentityType.ACI), + usernameHash, + Optional.empty()))); + + account.getUsernameHolds().forEach(usernameHold -> + constraintFutures.add(writeUsernameConstraint(account.getIdentifier(IdentityType.ACI), + usernameHold.usernameHash(), + Optional.of(Instant.ofEpochSecond(usernameHold.expirationSecs()))))); + + return CompletableFuture.allOf(constraintFutures.toArray(CompletableFuture[]::new)); + } + + private CompletableFuture writeConstraint( + final String tableName, + final UUID accountIdentifier, + final String keyName, + final AttributeValue keyValue) { + + return dynamoDbAsyncClient.putItem(PutItemRequest.builder() + .tableName(tableName) + .item(Map.of( + keyName, keyValue, + KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier))) + .build()) + .thenRun(Util.NOOP); + } + + private CompletableFuture writeUsernameConstraint( + final UUID accountIdentifier, + final byte[] usernameHash, + final Optional maybeExpiration) { + + final Map item = new HashMap<>(Map.of( + UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash), + UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier), + UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(maybeExpiration.isEmpty()) + )); + + maybeExpiration.ifPresent(expiration -> + item.put(UsernameTable.ATTR_TTL, AttributeValues.fromLong(expiration.getEpochSecond()))); + + return dynamoDbAsyncClient.putItem(PutItemRequest.builder() + .tableName(usernamesConstraintTableName) + .item(item) + .build()) + .thenRun(Util.NOOP); + } + @Nonnull private static String extractCancellationReasonCodes(final TransactionCanceledException exception) { return exception.cancellationReasons().stream() diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java index 374d44cc8..7d3d6fdda 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java @@ -75,6 +75,7 @@ import software.amazon.awssdk.services.dynamodb.DynamoDbClient; * Construct utilities commonly used by worker commands */ record CommandDependencies( + Accounts accounts, AccountsManager accountsManager, ProfilesManager profilesManager, ReportMessageManager reportMessageManager, @@ -290,6 +291,7 @@ record CommandDependencies( environment.lifecycle().manage(new ManagedAwsCrt()); return new CommandDependencies( + accounts, accountsManager, profilesManager, reportMessageManager, diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommand.java new file mode 100644 index 000000000..3e7d12b9b --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommand.java @@ -0,0 +1,84 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import com.google.common.annotations.VisibleForTesting; +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Metrics; +import java.time.Duration; +import net.sourceforge.argparse4j.inf.Subparser; +import org.whispersystems.textsecuregcm.metrics.MetricsUtil; +import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.Accounts; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.util.retry.Retry; + +public class RegenerateAccountConstraintDataCommand extends AbstractSinglePassCrawlAccountsCommand { + + @VisibleForTesting + static final String DRY_RUN_ARGUMENT = "dry-run"; + + @VisibleForTesting + static final String MAX_CONCURRENCY_ARGUMENT = "max-concurrency"; + + @VisibleForTesting + static final String RETRIES_ARGUMENT = "retries"; + + private static final String PROCESSED_ACCOUNTS_COUNTER_NAME = + MetricsUtil.name(RegenerateAccountConstraintDataCommand.class, "processedAccounts"); + + public RegenerateAccountConstraintDataCommand() { + super("regenerate-account-constraint-data", "Regenerates account constraint data from a core account table"); + } + + @Override + public void configure(final Subparser subparser) { + super.configure(subparser); + + subparser.addArgument("--dry-run") + .type(Boolean.class) + .dest(DRY_RUN_ARGUMENT) + .required(false) + .setDefault(true) + .help("If true, don’t actually write constraint data"); + + subparser.addArgument("--max-concurrency") + .type(Integer.class) + .dest(MAX_CONCURRENCY_ARGUMENT) + .setDefault(16) + .help("Max concurrency for DynamoDB operations"); + + subparser.addArgument("--retries") + .type(Integer.class) + .dest(RETRIES_ARGUMENT) + .setDefault(8) + .help("Maximum number of DynamoDB retries permitted per account"); + } + + @Override + protected void crawlAccounts(final Flux accountRecords) { + final boolean dryRun = getNamespace().getBoolean(DRY_RUN_ARGUMENT); + final int maxConcurrency = getNamespace().getInt(MAX_CONCURRENCY_ARGUMENT); + final int maxRetries = getNamespace().getInt(RETRIES_ARGUMENT); + + final Accounts accounts = getCommandDependencies().accounts(); + + final Counter processedAccountsCounter = Metrics.counter(PROCESSED_ACCOUNTS_COUNTER_NAME, + "dryRun", String.valueOf(dryRun)); + + accountRecords + .doOnNext(ignored -> processedAccountsCounter.increment()) + .flatMap(account -> dryRun + ? Mono.empty() + : Mono.fromFuture(() -> accounts.regenerateConstraints(account)) + .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(4)) + .onRetryExhaustedThrow((spec, rs) -> rs.failure())), + maxConcurrency) + .then() + .block(); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java index b7970505d..593e5c6bc 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java @@ -25,11 +25,13 @@ import com.fasterxml.jackson.core.JsonProcessingException; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.time.Instant; +import java.time.temporal.ChronoUnit; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -78,6 +80,7 @@ import software.amazon.awssdk.services.dynamodb.model.GetItemResponse; import software.amazon.awssdk.services.dynamodb.model.Put; import software.amazon.awssdk.services.dynamodb.model.PutItemRequest; import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure; +import software.amazon.awssdk.services.dynamodb.model.ScanRequest; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest; import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsResponse; @@ -118,6 +121,9 @@ class AccountsTest { private final TestClock clock = TestClock.pinned(Instant.EPOCH); private Accounts accounts; + private record UsernameConstraint(UUID accountIdentifier, boolean confirmed, Optional expiration) { + } + @BeforeEach void setupAccountsDao() { @@ -1685,7 +1691,7 @@ class AccountsTest { } @Test - public void testInvalidDeviceIdDeserialization() throws Exception { + void testInvalidDeviceIdDeserialization() throws Exception { final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID()); final Device device2 = generateDevice((byte) 64); account.addDevice(device2); @@ -1727,6 +1733,189 @@ class AccountsTest { assertInstanceOf(DeviceIdDeserializer.DeviceIdDeserializationException.class, cause); } + @Test + void testRegenerateConstraints() { + final Instant usernameHoldExpiration = clock.instant().plus(Accounts.USERNAME_HOLD_DURATION).truncatedTo(ChronoUnit.SECONDS); + + final Account account = nextRandomAccount(); + account.setUsernameHash(USERNAME_HASH_1); + account.setUsernameLinkDetails(UUID.randomUUID(), ENCRYPTED_USERNAME_1); + account.setUsernameHolds(List.of(new Account.UsernameHold(USERNAME_HASH_2, usernameHoldExpiration.getEpochSecond()))); + + writeAccountRecordWithoutConstraints(account); + accounts.regenerateConstraints(account).join(); + + // Check that constraints do what they should from a functional perspective + { + final Account conflictingNumberAccount = nextRandomAccount(); + conflictingNumberAccount.setNumber(account.getNumber(), account.getIdentifier(IdentityType.PNI)); + + assertThrows(AccountAlreadyExistsException.class, + () -> accounts.create(conflictingNumberAccount, Collections.emptyList())); + } + + { + final Account conflictingUsernameAccount = nextRandomAccount(); + createAccount(conflictingUsernameAccount); + + final CompletionException completionException = assertThrows(CompletionException.class, + () -> accounts.reserveUsernameHash(conflictingUsernameAccount, USERNAME_HASH_1, Accounts.USERNAME_HOLD_DURATION).join()); + + assertInstanceOf(UsernameHashNotAvailableException.class, completionException.getCause()); + } + + { + final Account conflictingUsernameHoldAccount = nextRandomAccount(); + createAccount(conflictingUsernameHoldAccount); + + final CompletionException completionException = assertThrows(CompletionException.class, + () -> accounts.reserveUsernameHash(conflictingUsernameHoldAccount, USERNAME_HASH_2, Accounts.USERNAME_HOLD_DURATION).join()); + + assertInstanceOf(UsernameHashNotAvailableException.class, completionException.getCause()); + } + + // Check that bare constraint records are written as expected + assertEquals(Optional.of(account.getIdentifier(IdentityType.ACI)), + getConstraintValue(Tables.NUMBERS.tableName(), Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()))); + + assertEquals(Optional.of(account.getIdentifier(IdentityType.ACI)), + getConstraintValue(Tables.PNI_ASSIGNMENTS.tableName(), Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getIdentifier(IdentityType.PNI)))); + + assertEquals(Optional.of(new UsernameConstraint(account.getIdentifier(IdentityType.ACI), true, Optional.empty())), + getUsernameConstraint(USERNAME_HASH_1)); + + assertEquals(Optional.of(new UsernameConstraint(account.getIdentifier(IdentityType.ACI), false, Optional.of(usernameHoldExpiration))), + getUsernameConstraint(USERNAME_HASH_2)); + } + + @Test + void testRegeneratedConstraintsMatchOriginalConstraints() { + final Instant usernameHoldExpiration = clock.instant().plus(Accounts.USERNAME_HOLD_DURATION).truncatedTo(ChronoUnit.SECONDS); + + final Account account = nextRandomAccount(); + account.setUsernameHash(USERNAME_HASH_1); + account.setUsernameLinkDetails(UUID.randomUUID(), ENCRYPTED_USERNAME_1); + account.setUsernameHolds(List.of(new Account.UsernameHold(USERNAME_HASH_2, usernameHoldExpiration.getEpochSecond()))); + + createAccount(account); + accounts.reserveUsernameHash(account, USERNAME_HASH_2, Accounts.USERNAME_HOLD_DURATION).join(); + accounts.confirmUsernameHash(account, USERNAME_HASH_2, ENCRYPTED_USERNAME_2).join(); + accounts.reserveUsernameHash(account, USERNAME_HASH_1, Accounts.USERNAME_HOLD_DURATION).join(); + accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join(); + + final Map originalE164ConstraintItem = + DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(Tables.NUMBERS.tableName()) + .key(Map.of(Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()))) + .build()) + .item(); + + final Map originalPniConstraintItem = + DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(Tables.PNI_ASSIGNMENTS.tableName()) + .key(Map.of(Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getIdentifier(IdentityType.PNI)))) + .build()) + .item(); + + final Set> originalUsernameConstraints = new HashSet<>( + DYNAMO_DB_EXTENSION.getDynamoDbClient().scan(ScanRequest.builder() + .tableName(Tables.USERNAMES.tableName()) + .build()) + .items()); + + accounts.delete(account.getIdentifier(IdentityType.ACI), Collections.emptyList()).join(); + + writeAccountRecordWithoutConstraints(account); + accounts.regenerateConstraints(account).join(); + + final Map regeneratedE164ConstraintItem = + DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(Tables.NUMBERS.tableName()) + .key(Map.of(Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()))) + .build()) + .item(); + + final Map regeneratedPniConstraintItem = + DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(Tables.PNI_ASSIGNMENTS.tableName()) + .key(Map.of(Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getIdentifier(IdentityType.PNI)))) + .build()) + .item(); + + final Set> regeneratedUsernameConstraints = new HashSet<>( + DYNAMO_DB_EXTENSION.getDynamoDbClient().scan(ScanRequest.builder() + .tableName(Tables.USERNAMES.tableName()) + .build()) + .items()); + + assertEquals(originalE164ConstraintItem, regeneratedE164ConstraintItem); + assertEquals(originalPniConstraintItem, regeneratedPniConstraintItem); + assertEquals(originalUsernameConstraints, regeneratedUsernameConstraints); + } + + private void writeAccountRecordWithoutConstraints(final Account account) { + final AttributeValue accountData; + + try { + accountData = AttributeValues.fromByteArray(Accounts.ACCOUNT_DDB_JSON_WRITER.writeValueAsBytes(account)); + } catch (final JsonProcessingException e) { + throw new IllegalArgumentException(e); + } + + final Map item = new HashMap<>(Map.of( + Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()), + Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()), + Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getPhoneNumberIdentifier()), + Accounts.ATTR_ACCOUNT_DATA, accountData, + Accounts.ATTR_VERSION, AttributeValues.fromInt(account.getVersion()), + Accounts.ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.isDiscoverableByPhoneNumber()))); + + account.getUnidentifiedAccessKey() + .map(AttributeValues::fromByteArray) + .ifPresent(uak -> item.put(Accounts.ATTR_UAK, uak)); + + DYNAMO_DB_EXTENSION.getDynamoDbClient().putItem(PutItemRequest.builder() + .tableName(Tables.ACCOUNTS.tableName()) + .item(item) + .build()); + } + + private Optional getConstraintValue(final String tableName, + final String keyName, + final AttributeValue keyValue) { + + final GetItemResponse response = DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(tableName) + .key(Map.of(keyName, keyValue)) + .build()); + + return response.hasItem() + ? Optional.ofNullable(AttributeValues.getUUID(response.item(), Accounts.KEY_ACCOUNT_UUID, null)) + : Optional.empty(); + } + + private Optional getUsernameConstraint(final byte[] usernameHash) { + final GetItemResponse response = DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder() + .tableName(Tables.USERNAMES.tableName()) + .key(Map.of(Accounts.UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash))) + .build()); + + if (response.hasItem()) { + final UUID accountIdentifier = + AttributeValues.getUUID(response.item(), Accounts.UsernameTable.ATTR_ACCOUNT_UUID, null); + + final boolean confirmed = AttributeValues.getBool(response.item(), Accounts.UsernameTable.ATTR_CONFIRMED, false); + + final Optional expiration = response.item().containsKey(Accounts.UsernameTable.ATTR_TTL) + ? Optional.of(Instant.ofEpochSecond(AttributeValues.getLong(response.item(), Accounts.UsernameTable.ATTR_TTL, 0))) + : Optional.empty(); + + return Optional.of(new UsernameConstraint(accountIdentifier, confirmed, expiration)); + } + + return Optional.empty(); + } + private static Device generateDevice(byte id) { return DevicesHelper.createDevice(id); } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java index 7605b2b76..4e01bd931 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java @@ -64,7 +64,8 @@ class FinishPushNotificationExperimentCommandTest { new PushNotificationExperimentSample<>(accountIdentifier, deviceId, true, "test", "test")); }); - commandDependencies = new CommandDependencies(accountsManager, + commandDependencies = new CommandDependencies(null, + accountsManager, null, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java index 70044e5bc..c2e5d028b 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java @@ -51,6 +51,7 @@ class NotifyIdleDevicesCommandTest { null, null, null, + null, messagesManager, null, null, diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommandTest.java new file mode 100644 index 000000000..544699623 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommandTest.java @@ -0,0 +1,96 @@ +/* + * Copyright 2025 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import net.sourceforge.argparse4j.inf.Namespace; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.Accounts; +import reactor.core.publisher.Flux; + +class RegenerateAccountConstraintDataCommandTest { + + private Accounts accounts; + + private static class TestRegenerateAccountConstraintDataCommand extends RegenerateAccountConstraintDataCommand { + + private final CommandDependencies commandDependencies; + private final Namespace namespace; + + TestRegenerateAccountConstraintDataCommand(final Accounts accounts, final boolean dryRun) { + commandDependencies = new CommandDependencies(accounts, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null); + + namespace = new Namespace(Map.of( + RegenerateAccountConstraintDataCommand.DRY_RUN_ARGUMENT, dryRun, + RegenerateAccountConstraintDataCommand.MAX_CONCURRENCY_ARGUMENT, 16, + RegenerateAccountConstraintDataCommand.RETRIES_ARGUMENT, 3)); + } + + @Override + protected CommandDependencies getCommandDependencies() { + return commandDependencies; + } + + @Override + protected Namespace getNamespace() { + return namespace; + } + } + + @BeforeEach + void setUp() { + accounts = mock(Accounts.class); + + when(accounts.regenerateConstraints(any())).thenReturn(CompletableFuture.completedFuture(null)); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void crawlAccounts(final boolean dryRun) { + final Account account = mock(Account.class); + + final RegenerateAccountConstraintDataCommand regenerateAccountConstraintDataCommand = + new TestRegenerateAccountConstraintDataCommand(accounts, dryRun); + + regenerateAccountConstraintDataCommand.crawlAccounts(Flux.just(account)); + + if (!dryRun) { + verify(accounts).regenerateConstraints(account); + } + + verifyNoMoreInteractions(accounts); + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java index 0a87de49b..471517680 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java @@ -64,6 +64,7 @@ class StartPushNotificationExperimentCommandTest { null, null, null, + null, pushNotificationExperimentSamples, null, null,