AggregatingLoadbalancerService implements LoadBalancerService { @Override public Completable removeLoadBalancer(RemoveLoadBalancerRequest removeLoadBalancerRequest, CallMetadata callMetadata) { JobId jobId = JobId.newBuilder().setId(removeLoadBalancerRequest.getJobId()).build(); final Observable<Empty> responseObservable = jobManagementServiceHelper.findJobInAllCells(jobId.getId(), callMetadata) .flatMap(response -> singleCellCall(response.getCell(), (client, responseObserver) -> client.removeLoadBalancer(removeLoadBalancerRequest, responseObserver), callMetadata)); return responseObservable.toCompletable(); } @Inject  AggregatingLoadbalancerService(CellConnector connector,
                                          CallMetadataResolver callMetadataResolver,
                                          GrpcConfiguration grpcConfiguration,
                                          AggregatingCellClient aggregatingClient,
                                          AggregatingJobManagementServiceHelper jobManagementServiceHelper); @Override Observable<GetAllLoadBalancersResult> getAllLoadBalancers(GetAllLoadBalancersRequest request, CallMetadata callMetadata); @Override Observable<GetJobLoadBalancersResult> getLoadBalancers(JobId jobId, CallMetadata callMetadata); @Override Completable addLoadBalancer(AddLoadBalancerRequest request, CallMetadata callMetadata); @Override Completable removeLoadBalancer(RemoveLoadBalancerRequest removeLoadBalancerRequest, CallMetadata callMetadata);  }
@Test public void removeLoadBalancer() { JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1); JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_2, LB_2); JobLoadBalancer jobLoadBalancer3 = new JobLoadBalancer(JOB_2, LB_3); final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(singletonList(jobLoadBalancer1)); final CellWithLoadBalancers cellWithLoadBalancersTwo = new CellWithLoadBalancers(new ArrayList<>(Arrays.asList(jobLoadBalancer2, jobLoadBalancer3))); final CellWithJobIds cellWithJobIdsOne = new CellWithJobIds(singletonList(JOB_1)); final CellWithJobIds cellWithJobIdsTwo = new CellWithJobIds(singletonList(JOB_2)); cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne); cellOne.getServiceRegistry().addService(cellWithJobIdsOne); cellTwo.getServiceRegistry().addService(cellWithLoadBalancersTwo); cellTwo.getServiceRegistry().addService(cellWithJobIdsTwo); final AssertableSubscriber<Void> resultSubscriber = service.removeLoadBalancer( RemoveLoadBalancerRequest.newBuilder().setJobId(JOB_2).setLoadBalancerId(LoadBalancerId.newBuilder().setId(LB_2).build()).build(), JUNIT_REST_CALL_METADATA ).test(); resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS); resultSubscriber.assertNoErrors(); resultSubscriber.assertNoValues(); resultSubscriber.assertCompleted(); final AssertableSubscriber<GetJobLoadBalancersResult> jobResults = service.getLoadBalancers( JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA ).test(); jobResults.awaitValueCount(1, 1, TimeUnit.SECONDS); jobResults.assertNoErrors(); final List<GetJobLoadBalancersResult> onNextEvents = jobResults.getOnNextEvents(); assertThat(onNextEvents.size()).isEqualTo(1); final List<LoadBalancerId> loadBalancersList = onNextEvents.get(0).getLoadBalancersList(); assertThat(loadBalancersList.size()).isEqualTo(1); assertThat(loadBalancersList.get(0).getId()).isEqualTo(LB_3); }