DRY gRPC tests, refactor error mapping

This commit is contained in:
Sergey Skrobotov
2023-09-08 16:08:59 -07:00
parent 29ca544c95
commit 977243ebfd
18 changed files with 637 additions and 525 deletions

View File

@@ -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) {