Add "additional specifiers" dimensions to open WebSocket counters and simplify meter construction

This commit is contained in:
Jon Chambers
2026-02-18 16:58:45 -05:00
committed by Jon Chambers
parent cb3363410c
commit f390aabb3a
7 changed files with 82 additions and 88 deletions

View File

@@ -6,7 +6,6 @@ import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
@@ -42,10 +41,10 @@ import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.filters.RemoteAddressFilter;
import org.whispersystems.textsecuregcm.push.ProvisioningManager;
import org.whispersystems.textsecuregcm.storage.ClientReleaseManager;
import org.whispersystems.textsecuregcm.tests.util.TestWebsocketListener;
import org.whispersystems.textsecuregcm.websocket.ProvisioningConnectListener;
import org.whispersystems.websocket.WebSocketResourceProviderFactory;
import org.whispersystems.websocket.WebsocketHeaders;
import org.whispersystems.websocket.configuration.WebSocketConfiguration;
import org.whispersystems.websocket.messages.InvalidMessageException;
import org.whispersystems.websocket.messages.WebSocketMessage;
@@ -110,7 +109,7 @@ public class ProvisioningTimeoutIntegrationTest {
.addFilter("RemoteAddressFilter", new RemoteAddressFilter())
.addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), false, "/*");
webSocketEnvironment.setConnectListener(
new ProvisioningConnectListener(mock(ProvisioningManager.class), scheduler, Duration.ofSeconds(5)));
new ProvisioningConnectListener(mock(ProvisioningManager.class), mock(ClientReleaseManager.class), scheduler, Duration.ofSeconds(5)));
final WebSocketResourceProviderFactory<AuthenticatedDevice> webSocketServlet =
new WebSocketResourceProviderFactory<>(webSocketEnvironment, AuthenticatedDevice.class,

View File

@@ -12,6 +12,7 @@ import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -23,9 +24,9 @@ import org.junit.jupiter.api.Test;
import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
import org.whispersystems.textsecuregcm.auth.DisconnectionRequestManager;
import org.whispersystems.textsecuregcm.identity.IdentityType;
import org.whispersystems.textsecuregcm.metrics.OpenWebSocketCounter;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.ClientReleaseManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.websocket.WebSocketClient;
import org.whispersystems.websocket.session.WebSocketSessionContext;
@@ -54,8 +55,8 @@ class AuthenticatedConnectListenerTest {
authenticatedConnectListener = new AuthenticatedConnectListener(accountsManager,
disconnectionRequestManager,
(_, _, _) -> authenticatedWebSocketConnection,
_ -> mock(OpenWebSocketCounter.class));
mock(ClientReleaseManager.class),
(_, _, _) -> authenticatedWebSocketConnection);
final Device device = mock(Device.class);
when(device.getId()).thenReturn(DEVICE_ID);
@@ -81,7 +82,8 @@ class AuthenticatedConnectListenerTest {
authenticatedConnectListener.onWebSocketConnect(webSocketSessionContext);
verify(disconnectionRequestManager).addListener(ACCOUNT_IDENTIFIER, DEVICE_ID, authenticatedWebSocketConnection);
verify(webSocketSessionContext).addWebsocketClosedListener(any());
// We expect one call from AuthenticatedConnectListener itself and one from OpenWebSocketCounter
verify(webSocketSessionContext, times(2)).addWebsocketClosedListener(any());
verify(authenticatedWebSocketConnection).start();
}
@@ -98,7 +100,8 @@ class AuthenticatedConnectListenerTest {
verify(webSocketClient).close(eq(1011), anyString());
verify(disconnectionRequestManager, never()).addListener(any(), anyByte(), any());
verify(webSocketSessionContext, never()).addWebsocketClosedListener(any());
// We expect one call from OpenWebSocketCounter, but none from AuthenticatedConnectListener itself
verify(webSocketSessionContext, times(1)).addWebsocketClosedListener(any());
verify(authenticatedWebSocketConnection, never()).start();
}
@@ -114,7 +117,8 @@ class AuthenticatedConnectListenerTest {
authenticatedConnectListener.onWebSocketConnect(webSocketSessionContext);
verify(disconnectionRequestManager).addListener(ACCOUNT_IDENTIFIER, DEVICE_ID, authenticatedWebSocketConnection);
verify(webSocketSessionContext).addWebsocketClosedListener(any());
// We expect one call from AuthenticatedConnectListener itself and one from OpenWebSocketCounter
verify(webSocketSessionContext, times(2)).addWebsocketClosedListener(any());
verify(authenticatedWebSocketConnection).start();
verify(webSocketClient).close(eq(1011), anyString());
@@ -125,7 +129,8 @@ class AuthenticatedConnectListenerTest {
authenticatedConnectListener.onWebSocketConnect(webSocketSessionContext);
verify(disconnectionRequestManager, never()).addListener(any(), anyByte(), any());
verify(webSocketSessionContext, never()).addWebsocketClosedListener(any());
// We expect one call from OpenWebSocketCounter, but none from AuthenticatedConnectListener itself
verify(webSocketSessionContext, times(1)).addWebsocketClosedListener(any());
verify(authenticatedWebSocketConnection, never()).start();
}
}

View File

@@ -11,7 +11,6 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import com.google.protobuf.InvalidProtocolBufferException;
import java.time.Duration;
@@ -24,6 +23,7 @@ import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.push.ProvisioningManager;
import org.whispersystems.textsecuregcm.storage.ClientReleaseManager;
import org.whispersystems.websocket.WebSocketClient;
import org.whispersystems.websocket.session.WebSocketSessionContext;
@@ -40,7 +40,7 @@ class ProvisioningConnectListenerTest {
provisioningManager = mock(ProvisioningManager.class);
scheduledExecutorService = mock(ScheduledExecutorService.class);
provisioningConnectListener =
new ProvisioningConnectListener(provisioningManager, scheduledExecutorService, TIMEOUT);
new ProvisioningConnectListener(provisioningManager, mock(ClientReleaseManager.class), scheduledExecutorService, TIMEOUT);
}
@Test