mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-20 21:28:02 +01:00
DRY gRPC tests, refactor error mapping
This commit is contained in:
@@ -12,12 +12,14 @@ import static org.mockito.Mockito.doThrow;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import org.apache.commons.lang3.RandomUtils;
|
||||
import org.mockito.Mockito;
|
||||
import org.whispersystems.textsecuregcm.configuration.secrets.SecretBytes;
|
||||
import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
|
||||
import org.whispersystems.textsecuregcm.limits.RateLimiter;
|
||||
import org.whispersystems.textsecuregcm.limits.RateLimiters;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
public final class MockUtils {
|
||||
|
||||
@@ -46,32 +48,80 @@ public final class MockUtils {
|
||||
}
|
||||
|
||||
public static void updateRateLimiterResponseToAllow(
|
||||
final RateLimiters rateLimitersMock,
|
||||
final RateLimiters.For handle,
|
||||
final RateLimiter mockRateLimiter,
|
||||
final String input) {
|
||||
final RateLimiter mockRateLimiter = Mockito.mock(RateLimiter.class);
|
||||
doReturn(mockRateLimiter).when(rateLimitersMock).forDescriptor(eq(handle));
|
||||
try {
|
||||
doNothing().when(mockRateLimiter).validate(eq(input));
|
||||
doReturn(CompletableFuture.completedFuture(null)).when(mockRateLimiter).validateAsync(eq(input));
|
||||
doReturn(Mono.fromFuture(CompletableFuture.completedFuture(null))).when(mockRateLimiter).validateReactive(eq(input));
|
||||
} catch (final RateLimitExceededException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public static void updateRateLimiterResponseToAllow(
|
||||
final RateLimiter mockRateLimiter,
|
||||
final UUID input) {
|
||||
try {
|
||||
doNothing().when(mockRateLimiter).validate(eq(input));
|
||||
doReturn(CompletableFuture.completedFuture(null)).when(mockRateLimiter).validateAsync(eq(input));
|
||||
doReturn(Mono.fromFuture(CompletableFuture.completedFuture(null))).when(mockRateLimiter).validateReactive(eq(input));
|
||||
} catch (final RateLimitExceededException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public static void updateRateLimiterResponseToAllow(
|
||||
final RateLimiters rateLimitersMock,
|
||||
final RateLimiters.For handle,
|
||||
final String input) {
|
||||
final RateLimiter mockRateLimiter = Mockito.mock(RateLimiter.class);
|
||||
doReturn(mockRateLimiter).when(rateLimitersMock).forDescriptor(eq(handle));
|
||||
updateRateLimiterResponseToAllow(mockRateLimiter, input);
|
||||
}
|
||||
|
||||
public static void updateRateLimiterResponseToAllow(
|
||||
final RateLimiters rateLimitersMock,
|
||||
final RateLimiters.For handle,
|
||||
final UUID input) {
|
||||
final RateLimiter mockRateLimiter = Mockito.mock(RateLimiter.class);
|
||||
doReturn(mockRateLimiter).when(rateLimitersMock).forDescriptor(eq(handle));
|
||||
updateRateLimiterResponseToAllow(mockRateLimiter, input);
|
||||
}
|
||||
|
||||
public static Duration updateRateLimiterResponseToFail(
|
||||
final RateLimiter mockRateLimiter,
|
||||
final String input,
|
||||
final Duration retryAfter,
|
||||
final boolean legacyStatusCode) {
|
||||
try {
|
||||
doNothing().when(mockRateLimiter).validate(eq(input));
|
||||
final RateLimitExceededException exception = new RateLimitExceededException(retryAfter, legacyStatusCode);
|
||||
doThrow(exception).when(mockRateLimiter).validate(eq(input));
|
||||
doReturn(CompletableFuture.failedFuture(exception)).when(mockRateLimiter).validateAsync(eq(input));
|
||||
doReturn(Mono.fromFuture(CompletableFuture.failedFuture(exception))).when(mockRateLimiter).validateReactive(eq(input));
|
||||
return retryAfter;
|
||||
} catch (final RateLimitExceededException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public static void updateRateLimiterResponseToFail(
|
||||
public static Duration updateRateLimiterResponseToFail(
|
||||
final RateLimiter mockRateLimiter,
|
||||
final UUID input,
|
||||
final Duration retryAfter,
|
||||
final boolean legacyStatusCode) {
|
||||
try {
|
||||
final RateLimitExceededException exception = new RateLimitExceededException(retryAfter, legacyStatusCode);
|
||||
doThrow(exception).when(mockRateLimiter).validate(eq(input));
|
||||
doReturn(CompletableFuture.failedFuture(exception)).when(mockRateLimiter).validateAsync(eq(input));
|
||||
doReturn(Mono.fromFuture(CompletableFuture.failedFuture(exception))).when(mockRateLimiter).validateReactive(eq(input));
|
||||
return retryAfter;
|
||||
} catch (final RateLimitExceededException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public static Duration updateRateLimiterResponseToFail(
|
||||
final RateLimiters rateLimitersMock,
|
||||
final RateLimiters.For handle,
|
||||
final String input,
|
||||
@@ -79,14 +129,10 @@ public final class MockUtils {
|
||||
final boolean legacyStatusCode) {
|
||||
final RateLimiter mockRateLimiter = Mockito.mock(RateLimiter.class);
|
||||
doReturn(mockRateLimiter).when(rateLimitersMock).forDescriptor(eq(handle));
|
||||
try {
|
||||
doThrow(new RateLimitExceededException(retryAfter, legacyStatusCode)).when(mockRateLimiter).validate(eq(input));
|
||||
} catch (final RateLimitExceededException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
return updateRateLimiterResponseToFail(mockRateLimiter, input, retryAfter, legacyStatusCode);
|
||||
}
|
||||
|
||||
public static void updateRateLimiterResponseToFail(
|
||||
public static Duration updateRateLimiterResponseToFail(
|
||||
final RateLimiters rateLimitersMock,
|
||||
final RateLimiters.For handle,
|
||||
final UUID input,
|
||||
@@ -94,11 +140,7 @@ public final class MockUtils {
|
||||
final boolean legacyStatusCode) {
|
||||
final RateLimiter mockRateLimiter = Mockito.mock(RateLimiter.class);
|
||||
doReturn(mockRateLimiter).when(rateLimitersMock).forDescriptor(eq(handle));
|
||||
try {
|
||||
doThrow(new RateLimitExceededException(retryAfter, legacyStatusCode)).when(mockRateLimiter).validate(eq(input));
|
||||
} catch (final RateLimitExceededException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
return updateRateLimiterResponseToFail(mockRateLimiter, input, retryAfter, legacyStatusCode);
|
||||
}
|
||||
|
||||
public static SecretBytes randomSecretBytes(final int size) {
|
||||
|
||||
Reference in New Issue
Block a user