ServingServiceGRpcController extends ServingServiceImplBase { @Override public void getOnlineFeatures( GetOnlineFeaturesRequest request, StreamObserver<GetOnlineFeaturesResponse> responseObserver) { Span span = tracer.buildSpan("getOnlineFeatures").start(); try (Scope scope = tracer.scopeManager().activate(span, false)) { if (request.getProject() != null && !request.getProject().isEmpty()) { this.authorizationService.authorizeRequest( SecurityContextHolder.getContext(), request.getProject()); } else { this.checkProjectAccess(request.getFeaturesList()); } RequestHelper.validateOnlineRequest(request); GetOnlineFeaturesResponse onlineFeatures = servingService.getOnlineFeatures(request); responseObserver.onNext(onlineFeatures); responseObserver.onCompleted(); } catch (SpecRetrievalException e) { log.error("Failed to retrieve specs in SpecService", e); responseObserver.onError( Status.NOT_FOUND.withDescription(e.getMessage()).withCause(e).asException()); } catch (AccessDeniedException e) { log.info(String.format("User prevented from accessing one of the projects in request")); responseObserver.onError( Status.PERMISSION_DENIED .withDescription(e.getMessage()) .withCause(e) .asRuntimeException()); } catch (Exception e) { log.warn("Failed to get Online Features", e); responseObserver.onError(e); } span.finish(); } @Autowired  ServingServiceGRpcController(
      AuthorizationService authorizationService,
      ServingService servingService,
      FeastProperties feastProperties,
      Tracer tracer); @Override void getFeastServingInfo(
      GetFeastServingInfoRequest request,
      StreamObserver<GetFeastServingInfoResponse> responseObserver); @Override void getOnlineFeatures(
      GetOnlineFeaturesRequest request,
      StreamObserver<GetOnlineFeaturesResponse> responseObserver); @Override void getBatchFeatures(
      GetBatchFeaturesRequest request, StreamObserver<GetBatchFeaturesResponse> responseObserver); @Override void getJob(GetJobRequest request, StreamObserver<GetJobResponse> responseObserver);  }
@Test public void shouldPassValidRequestAsIs() { service = getServingServiceGRpcController(false); service.getOnlineFeatures(validRequest, mockStreamObserver); Mockito.verify(mockServingService).getOnlineFeatures(validRequest); }
@Test public void shouldCallOnErrorIfEntityDatasetIsNotSet() { service = getServingServiceGRpcController(false); GetOnlineFeaturesRequest missingEntityName = GetOnlineFeaturesRequest.newBuilder(validRequest).clearEntityRows().build(); service.getOnlineFeatures(missingEntityName, mockStreamObserver); Mockito.verify(mockStreamObserver).onError(Mockito.any(StatusRuntimeException.class)); }
@Test public void shouldPassValidRequestAsIsIfRequestIsAuthorized() { service = getServingServiceGRpcController(true); SecurityContext context = mock(SecurityContext.class); SecurityContextHolder.setContext(context); when(context.getAuthentication()).thenReturn(authentication); doReturn(AuthorizationResult.success()) .when(authProvider) .checkAccessToProject(anyString(), any(Authentication.class)); service.getOnlineFeatures(validRequest, mockStreamObserver); Mockito.verify(mockServingService).getOnlineFeatures(validRequest); }
@Test public void shouldThrowErrorOnValidRequestIfRequestIsUnauthorized() { service = getServingServiceGRpcController(true); SecurityContext context = mock(SecurityContext.class); SecurityContextHolder.setContext(context); when(context.getAuthentication()).thenReturn(authentication); doReturn(AuthorizationResult.failed(null)) .when(authProvider) .checkAccessToProject(anyString(), any(Authentication.class)); service.getOnlineFeatures(validRequest, mockStreamObserver); Mockito.verify(mockStreamObserver).onError(Mockito.any(StatusRuntimeException.class)); }