Add a remote address interceptor to base gRPC tests

This commit is contained in:
Jon Chambers 2023-10-04 15:16:35 -04:00 committed by Jon Chambers
parent f55504c665
commit eaa868cf06
3 changed files with 16 additions and 10 deletions

View File

@ -26,23 +26,24 @@ public final class GrpcTestUtils {
// noop // noop
} }
public static MockAuthenticationInterceptor setupAuthenticatedExtension( public static void setupAuthenticatedExtension(
final GrpcServerExtension extension, final GrpcServerExtension extension,
final MockAuthenticationInterceptor mockAuthenticationInterceptor,
final MockRemoteAddressInterceptor mockRemoteAddressInterceptor,
final UUID authenticatedAci, final UUID authenticatedAci,
final long authenticatedDeviceId, final long authenticatedDeviceId,
final BindableService service) { final BindableService service) {
final MockAuthenticationInterceptor mockAuthenticationInterceptor = new MockAuthenticationInterceptor();
mockAuthenticationInterceptor.setAuthenticatedDevice(authenticatedAci, authenticatedDeviceId); mockAuthenticationInterceptor.setAuthenticatedDevice(authenticatedAci, authenticatedDeviceId);
extension.getServiceRegistry() extension.getServiceRegistry()
.addService(ServerInterceptors.intercept(service, mockAuthenticationInterceptor, new ErrorMappingInterceptor())); .addService(ServerInterceptors.intercept(service, mockRemoteAddressInterceptor, mockAuthenticationInterceptor, new ErrorMappingInterceptor()));
return mockAuthenticationInterceptor;
} }
public static void setupUnauthenticatedExtension( public static void setupUnauthenticatedExtension(
final GrpcServerExtension extension, final GrpcServerExtension extension,
final MockRemoteAddressInterceptor mockRemoteAddressInterceptor,
final BindableService service) { final BindableService service) {
extension.getServiceRegistry() extension.getServiceRegistry()
.addService(ServerInterceptors.intercept(service, new ErrorMappingInterceptor())); .addService(ServerInterceptors.intercept(service, mockRemoteAddressInterceptor, new ErrorMappingInterceptor()));
} }
public static void assertStatusException(final Status expected, final Executable serviceCall) { public static void assertStatusException(final Status expected, final Executable serviceCall) {

View File

@ -28,7 +28,7 @@ public class MockRemoteAddressInterceptor implements ServerInterceptor {
final Metadata headers, final Metadata headers,
final ServerCallHandler<ReqT, RespT> next) { final ServerCallHandler<ReqT, RespT> next) {
return remoteAddress != null return remoteAddress == null
? next.startCall(serverCall, headers) ? next.startCall(serverCall, headers)
: Contexts.interceptCall( : Contexts.interceptCall(
Context.current().withValue(RemoteAddressUtil.REMOTE_ADDRESS_CONTEXT_KEY, remoteAddress), Context.current().withValue(RemoteAddressUtil.REMOTE_ADDRESS_CONTEXT_KEY, remoteAddress),

View File

@ -59,7 +59,8 @@ public abstract class SimpleBaseGrpcTest<SERVICE extends BindableService, STUB e
private AutoCloseable mocksCloseable; private AutoCloseable mocksCloseable;
private MockAuthenticationInterceptor mockAuthenticationInterceptor; private final MockAuthenticationInterceptor mockAuthenticationInterceptor = new MockAuthenticationInterceptor();
private final MockRemoteAddressInterceptor mockRemoteAddressInterceptor = new MockRemoteAddressInterceptor();
private SERVICE service; private SERVICE service;
@ -112,9 +113,9 @@ public abstract class SimpleBaseGrpcTest<SERVICE extends BindableService, STUB e
protected void baseSetup() { protected void baseSetup() {
mocksCloseable = MockitoAnnotations.openMocks(this); mocksCloseable = MockitoAnnotations.openMocks(this);
service = requireNonNull(createServiceBeforeEachTest(), "created service must not be `null`"); service = requireNonNull(createServiceBeforeEachTest(), "created service must not be `null`");
mockAuthenticationInterceptor = GrpcTestUtils.setupAuthenticatedExtension( GrpcTestUtils.setupAuthenticatedExtension(
GRPC_SERVER_EXTENSION_AUTHENTICATED, AUTHENTICATED_ACI, AUTHENTICATED_DEVICE_ID, service); GRPC_SERVER_EXTENSION_AUTHENTICATED, mockAuthenticationInterceptor, mockRemoteAddressInterceptor, AUTHENTICATED_ACI, AUTHENTICATED_DEVICE_ID, service);
GrpcTestUtils.setupUnauthenticatedExtension(GRPC_SERVER_EXTENSION_UNAUTHENTICATED, service); GrpcTestUtils.setupUnauthenticatedExtension(GRPC_SERVER_EXTENSION_UNAUTHENTICATED, mockRemoteAddressInterceptor, service);
try { try {
authenticatedServiceStub = createStub(GRPC_SERVER_EXTENSION_AUTHENTICATED.getChannel()); authenticatedServiceStub = createStub(GRPC_SERVER_EXTENSION_AUTHENTICATED.getChannel());
unauthenticatedServiceStub = createStub(GRPC_SERVER_EXTENSION_UNAUTHENTICATED.getChannel()); unauthenticatedServiceStub = createStub(GRPC_SERVER_EXTENSION_UNAUTHENTICATED.getChannel());
@ -143,4 +144,8 @@ public abstract class SimpleBaseGrpcTest<SERVICE extends BindableService, STUB e
protected STUB unauthenticatedServiceStub() { protected STUB unauthenticatedServiceStub() {
return unauthenticatedServiceStub; return unauthenticatedServiceStub;
} }
protected MockRemoteAddressInterceptor getMockRemoteAddressInterceptor() {
return mockRemoteAddressInterceptor;
}
} }