Send directory updates in batches.

This commit is contained in:
Jon Chambers
2021-05-03 12:33:31 -04:00
committed by Jon Chambers
parent 30c9968928
commit 8fdbcbef44
5 changed files with 121 additions and 44 deletions

View File

@@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.sqs;
import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.model.MessageAttributeValue;
import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
import com.amazonaws.services.sqs.model.SendMessageRequest;
import junitparams.JUnitParamsRunner;
import junitparams.Parameters;
@@ -43,13 +44,50 @@ public class DirectoryQueueTest {
directoryQueue.refreshRegisteredUser(account);
final ArgumentCaptor<SendMessageRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
verify(sqs).sendMessage(requestCaptor.capture());
final ArgumentCaptor<SendMessageBatchRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
verify(sqs).sendMessageBatch(requestCaptor.capture());
final Map<String, MessageAttributeValue> messageAttributes = requestCaptor.getValue().getMessageAttributes();
assertEquals(1, requestCaptor.getValue().getEntries().size());
final Map<String, MessageAttributeValue> messageAttributes = requestCaptor.getValue().getEntries().get(0).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(expectedAction), messageAttributes.get("action"));
}
@Test
public void testRefreshBatch() {
final AmazonSQS sqs = mock(AmazonSQS.class);
final DirectoryQueue directoryQueue = new DirectoryQueue(List.of("sqs://test"), sqs);
final Account discoverableAccount = mock(Account.class);
when(discoverableAccount.getNumber()).thenReturn("+18005556543");
when(discoverableAccount.getUuid()).thenReturn(UUID.randomUUID());
when(discoverableAccount.isEnabled()).thenReturn(true);
when(discoverableAccount.isDiscoverableByPhoneNumber()).thenReturn(true);
final Account undiscoverableAccount = mock(Account.class);
when(undiscoverableAccount.getNumber()).thenReturn("+18005550987");
when(undiscoverableAccount.getUuid()).thenReturn(UUID.randomUUID());
when(undiscoverableAccount.isEnabled()).thenReturn(true);
when(undiscoverableAccount.isDiscoverableByPhoneNumber()).thenReturn(false);
directoryQueue.refreshRegisteredUsers(List.of(discoverableAccount, undiscoverableAccount));
final ArgumentCaptor<SendMessageBatchRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
verify(sqs).sendMessageBatch(requestCaptor.capture());
assertEquals(2, requestCaptor.getValue().getEntries().size());
final Map<String, MessageAttributeValue> discoverableAccountAttributes = requestCaptor.getValue().getEntries().get(0).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(discoverableAccount.getNumber()), discoverableAccountAttributes.get("id"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(discoverableAccount.getUuid().toString()), discoverableAccountAttributes.get("uuid"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue("add"), discoverableAccountAttributes.get("action"));
final Map<String, MessageAttributeValue> undiscoverableAccountAttributes = requestCaptor.getValue().getEntries().get(1).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(undiscoverableAccount.getNumber()), undiscoverableAccountAttributes.get("id"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(undiscoverableAccount.getUuid().toString()), undiscoverableAccountAttributes.get("uuid"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue("delete"), undiscoverableAccountAttributes.get("action"));
}
@Test
public void testSendMessageMultipleQueues() {
final AmazonSQS sqs = mock(AmazonSQS.class);
@@ -63,11 +101,13 @@ public class DirectoryQueueTest {
directoryQueue.refreshRegisteredUser(account);
final ArgumentCaptor<SendMessageRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
verify(sqs, times(2)).sendMessage(requestCaptor.capture());
final ArgumentCaptor<SendMessageBatchRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
verify(sqs, times(2)).sendMessageBatch(requestCaptor.capture());
for (final SendMessageRequest sendMessageRequest : requestCaptor.getAllValues()) {
final Map<String, MessageAttributeValue> messageAttributes = sendMessageRequest.getMessageAttributes();
for (final SendMessageBatchRequest sendMessageBatchRequest : requestCaptor.getAllValues()) {
assertEquals(1, requestCaptor.getValue().getEntries().size());
final Map<String, MessageAttributeValue> messageAttributes = sendMessageBatchRequest.getEntries().get(0).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue("add"), messageAttributes.get("action"));
}
}

View File

@@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.tests.storage;
import org.junit.Before;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerRestartException;
@@ -22,6 +23,7 @@ import java.util.Set;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.*;
public class PushFeedbackProcessorTest {
@@ -131,8 +133,10 @@ public class PushFeedbackProcessorTest {
verify(accountsManager).update(eq(stillActiveAccount));
verify(directoryQueue).refreshRegisteredUser(undiscoverableAccount);
verify(directoryQueue).refreshRegisteredUser(uninstalledAccount);
final ArgumentCaptor<List<Account>> refreshedAccountArgumentCaptor = ArgumentCaptor.forClass(List.class);
verify(directoryQueue).refreshRegisteredUsers(refreshedAccountArgumentCaptor.capture());
assertTrue(refreshedAccountArgumentCaptor.getValue().containsAll(List.of(undiscoverableAccount, uninstalledAccount)));
}