From 0585f862cbac34c7525b583397d7aafc3e8eab49 Mon Sep 17 00:00:00 2001 From: Chris Eager Date: Fri, 11 Apr 2025 15:34:12 -0500 Subject: [PATCH] Add regression test for set profile badges calculation --- .../controllers/ProfileControllerTest.java | 52 ++++++++++++++ .../grpc/ProfileGrpcServiceTest.java | 47 ++++++++++++- .../tests/util/AccountsHelper.java | 67 ++++++++++++++++++- 3 files changed, 164 insertions(+), 2 deletions(-) diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ProfileControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ProfileControllerTest.java index c3f61d9a6..1c3aaee84 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ProfileControllerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/ProfileControllerTest.java @@ -1140,6 +1140,58 @@ class ProfileControllerTest { new AccountBadge("TEST2", Instant.ofEpochSecond(42 + 86400), false), new AccountBadge("TEST3", Instant.ofEpochSecond(42 + 86400), false)); } + + } + + @Test + void testSetProfileBadgeAfterUpdateTries() throws Exception { + final ProfileKeyCommitment commitment = new ProfileKey(new byte[32]).getCommitment( + new ServiceId.Aci(AuthHelper.VALID_UUID)); + + final byte[] name = TestRandomUtil.nextBytes(81); + final byte[] emoji = TestRandomUtil.nextBytes(60); + final byte[] about = TestRandomUtil.nextBytes(156); + final String version = versionHex("anotherversion"); + + clearInvocations(AuthHelper.VALID_ACCOUNT_TWO); + reset(accountsManager); + final int accountsManagerUpdateRetryCount = 2; + AccountsHelper.setupMockUpdateWithRetries(accountsManager, accountsManagerUpdateRetryCount); + // set up two invocations -- one for each AccountsManager#update try + when(AuthHelper.VALID_ACCOUNT_TWO.getBadges()) + .thenReturn(List.of( + new AccountBadge("TEST2", Instant.ofEpochSecond(42 + 86400), true), + new AccountBadge("TEST3", Instant.ofEpochSecond(42 + 86400), true) + )) + .thenReturn(List.of( + new AccountBadge("TEST2", Instant.ofEpochSecond(42 + 86400), true), + new AccountBadge("TEST3", Instant.ofEpochSecond(42 + 86400), true), + new AccountBadge("TEST4", Instant.ofEpochSecond(43 + 86400), true) + )); + + try (final Response response = resources.getJerseyTest() + .target("/v1/profile/") + .request() + .header("Authorization", AuthHelper.getAuthHeader(AuthHelper.VALID_UUID_TWO, AuthHelper.VALID_PASSWORD_TWO)) + .put(Entity.entity(new CreateProfileRequest(commitment, version, name, emoji, about, null, false, false, + Optional.of(List.of("TEST1")), null), MediaType.APPLICATION_JSON_TYPE))) { + + assertThat(response.getStatus()).isEqualTo(200); + assertThat(response.hasEntity()).isFalse(); + + //noinspection unchecked + final ArgumentCaptor> badgeCaptor = ArgumentCaptor.forClass(List.class); + verify(AuthHelper.VALID_ACCOUNT_TWO, times(accountsManagerUpdateRetryCount)).setBadges(refEq(clock), badgeCaptor.capture()); + // since the stubbing of getBadges() is brittle, we need to verify the number of invocations, to protect against upstream changes + verify(AuthHelper.VALID_ACCOUNT_TWO, times(accountsManagerUpdateRetryCount)).getBadges(); + + final List badges = badgeCaptor.getValue(); + assertThat(badges).isNotNull().hasSize(4).containsOnly( + new AccountBadge("TEST1", Instant.ofEpochSecond(42 + 86400), true), + new AccountBadge("TEST2", Instant.ofEpochSecond(42 + 86400), false), + new AccountBadge("TEST3", Instant.ofEpochSecond(42 + 86400), false), + new AccountBadge("TEST4", Instant.ofEpochSecond(43 + 86400), false)); + } } @ParameterizedTest diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ProfileGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ProfileGrpcServiceTest.java index e6ab71548..3073a8e5a 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ProfileGrpcServiceTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ProfileGrpcServiceTest.java @@ -15,7 +15,9 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.refEq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; @@ -30,6 +32,7 @@ import java.nio.charset.StandardCharsets; import java.time.Clock; import java.time.Duration; import java.time.Instant; +import java.time.ZoneId; import java.time.temporal.ChronoUnit; import java.util.Collections; import java.util.List; @@ -93,11 +96,13 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters; import org.whispersystems.textsecuregcm.s3.PolicySigner; import org.whispersystems.textsecuregcm.s3.PostPolicyGenerator; import org.whispersystems.textsecuregcm.storage.Account; +import org.whispersystems.textsecuregcm.storage.AccountBadge; import org.whispersystems.textsecuregcm.storage.AccountsManager; import org.whispersystems.textsecuregcm.storage.DeviceCapability; import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager; import org.whispersystems.textsecuregcm.storage.ProfilesManager; import org.whispersystems.textsecuregcm.storage.VersionedProfile; +import org.whispersystems.textsecuregcm.tests.util.AccountsHelper; import org.whispersystems.textsecuregcm.tests.util.AuthHelper; import org.whispersystems.textsecuregcm.tests.util.ProfileTestHelper; import org.whispersystems.textsecuregcm.util.MockUtils; @@ -144,6 +149,8 @@ public class ProfileGrpcServiceTest extends SimpleBaseGrpcTest dynamicConfigurationManager = mock(DynamicConfigurationManager.class); @@ -203,8 +210,10 @@ public class ProfileGrpcServiceTest extends SimpleBaseGrpcTest> badgeCaptor = ArgumentCaptor.forClass(List.class); + verify(account, times(2)).setBadges(refEq(clock), badgeCaptor.capture()); + // since the stubbing of getBadges() is brittle, we need to verify the number of invocations, to protect against upstream changes + verify(account, times(accountsManagerUpdateRetryCount)).getBadges(); + + assertEquals(List.of( + new AccountBadge("TEST3", Instant.ofEpochSecond(41), true), + new AccountBadge("TEST2", Instant.ofEpochSecond(41), false)), + badgeCaptor.getValue()); + } + @ParameterizedTest @EnumSource(value = org.signal.chat.common.IdentityType.class, names = {"IDENTITY_TYPE_ACI", "IDENTITY_TYPE_PNI"}) void getUnversionedProfile(final IdentityType identityType) { diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java index e7be0aac6..e20bf0651 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java @@ -31,8 +31,8 @@ import org.whispersystems.textsecuregcm.auth.SaltedTokenHash; import org.whispersystems.textsecuregcm.entities.AccountAttributes; import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.AccountsManager; -import org.whispersystems.textsecuregcm.storage.DeviceSpec; import org.whispersystems.textsecuregcm.storage.Device; +import org.whispersystems.textsecuregcm.storage.DeviceSpec; import org.whispersystems.textsecuregcm.util.SystemMapper; public class AccountsHelper { @@ -62,6 +62,71 @@ public class AccountsHelper { setupMockUpdate(mockAccountsManager, false); } + /** + * Sets up stubbing for: + *
    + *
  • {@link AccountsManager#update(Account, Consumer)}
    • + *
  • {@link AccountsManager#updateAsync(Account, Consumer)}
    • + *
  • {@link AccountsManager#updateDevice(Account, byte, Consumer)}
    • + *
  • {@link AccountsManager#updateDeviceAsync(Account, byte, Consumer)}
    • + *
+ * + * with multiple calls to the {@link Consumer}. This simulates retries from {@link org.whispersystems.textsecuregcm.storage.ContestedOptimisticLockException}. + * Callers will typically set up stubbing for relevant {@link Account} methods with multiple {@link org.mockito.stubbing.OngoingStubbing#thenReturn(Object)} + * calls: + * 
+   *   // example stubbing
+   *   when(account.getNextDeviceId())
+   *     .thenReturn(2)
+   *     .thenReturn(3);
+   *
+ */ + @SuppressWarnings("unchecked") + public static void setupMockUpdateWithRetries(final AccountsManager mockAccountsManager, final int retryCount) { + when(mockAccountsManager.update(any(), any())).thenAnswer(answer -> { + final Account account = answer.getArgument(0, Account.class); + + for (int i = 0; i < retryCount; i++) { + answer.getArgument(1, Consumer.class).accept(account); + } + + return copyAndMarkStale(account); + }); + + when(mockAccountsManager.updateAsync(any(), any())).thenAnswer(answer -> { + final Account account = answer.getArgument(0, Account.class); + + for (int i = 0; i < retryCount; i++) { + answer.getArgument(1, Consumer.class).accept(account); + } + + return CompletableFuture.completedFuture(copyAndMarkStale(account)); + }); + + when(mockAccountsManager.updateDevice(any(), anyByte(), any())).thenAnswer(answer -> { + final Account account = answer.getArgument(0, Account.class); + final byte deviceId = answer.getArgument(1, Byte.class); + + for (int i = 0; i < retryCount; i++) { + account.getDevice(deviceId).ifPresent(answer.getArgument(2, Consumer.class)); + } + + return copyAndMarkStale(account); + }); + + when(mockAccountsManager.updateDeviceAsync(any(), anyByte(), any())).thenAnswer(answer -> { + final Account account = answer.getArgument(0, Account.class); + final byte deviceId = answer.getArgument(1, Byte.class); + + for (int i = 0; i < retryCount; i++) { + account.getDevice(deviceId).ifPresent(answer.getArgument(2, Consumer.class)); + } + + return CompletableFuture.completedFuture(copyAndMarkStale(account)); + }); + } + + @SuppressWarnings("unchecked") private static void setupMockUpdate(final AccountsManager mockAccountsManager, final boolean markStale) { when(mockAccountsManager.update(any(), any())).thenAnswer(answer -> { final Account account = answer.getArgument(0, Account.class);