AggregatingJobServiceGateway implements JobServiceGateway { @Override public Completable updateJobCapacity(JobCapacityUpdate request, CallMetadata callMetadata) { Observable<Empty> result = jobManagementServiceHelper.findJobInAllCells(request.getJobId(), callMetadata) .flatMap(response -> singleCellCall(response.getCell(), (client, streamObserver) -> client.updateJobCapacity(request, streamObserver), callMetadata) ); return result.toCompletable(); } @Inject  AggregatingJobServiceGateway(GrpcConfiguration grpcConfiguration,
                                        TitusFederationConfiguration federationConfiguration,
                                        CellConnector connector,
                                        CellRouter router,
                                        AggregatingCellClient aggregatingClient,
                                        AggregatingJobManagementServiceHelper jobManagementServiceHelper); @Override Observable<String> createJob(JobDescriptor jobDescriptor, CallMetadata callMetadata); @Override Completable updateJobCapacity(JobCapacityUpdate request, CallMetadata callMetadata); @Override Completable updateJobCapacityWithOptionalAttributes(JobCapacityUpdateWithOptionalAttributes request, CallMetadata callMetadata); @Override Completable updateJobProcesses(JobProcessesUpdate request, CallMetadata callMetadata); @Override Completable updateJobStatus(JobStatusUpdate request, CallMetadata callMetadata); @Override Mono<Void> updateJobDisruptionBudget(JobDisruptionBudgetUpdate request, CallMetadata callMetadata); @Override Mono<Void> updateJobAttributes(JobAttributesUpdate request, CallMetadata callMetadata); @Override Mono<Void> deleteJobAttributes(JobAttributesDeleteRequest request, CallMetadata callMetadata); @Override Observable<Job> findJob(String jobId, CallMetadata callMetadata); @Override Observable<JobQueryResult> findJobs(JobQuery request, CallMetadata callMetadata); @Override Observable<JobChangeNotification> observeJob(String jobId, CallMetadata callMetadata); @Override Observable<JobChangeNotification> observeJobs(ObserveJobsQuery query, CallMetadata callMetadata); @Override Completable killJob(String jobId, CallMetadata callMetadata); @Override Observable<Task> findTask(String taskId, CallMetadata callMetadata); @Override Observable<TaskQueryResult> findTasks(TaskQuery request, CallMetadata callMetadata); @Override Completable killTask(TaskKillRequest request, CallMetadata callMetadata); @Override Completable updateTaskAttributes(TaskAttributesUpdate attributesUpdate, CallMetadata callMetadata); @Override Completable deleteTaskAttributes(TaskAttributesDeleteRequest deleteRequest, CallMetadata callMetadata); @Override Completable moveTask(TaskMoveRequest taskMoveRequest, CallMetadata callMetadata);  }
@Test public void singleJobUpdates() { Random random = new Random(); List<String> cellOneSnapshot = new ArrayList<>(dataGenerator.newServiceJobs(12)).stream() .map(job -> job.getId()) .collect(Collectors.toList()); List<String> cellTwoSnapshot = new ArrayList<>(dataGenerator.newBatchJobs(7)).stream() .map(job -> job.getId()) .collect(Collectors.toList()); CellWithJobIds cellOneService = new CellWithJobIds(cellOneSnapshot); CellWithJobIds cellTwoService = new CellWithJobIds(cellTwoSnapshot); cellOne.getServiceRegistry().addService(cellOneService); cellTwo.getServiceRegistry().addService(cellTwoService); String cellOneJobId = cellOneSnapshot.get(random.nextInt(cellOneSnapshot.size())); String cellTwoJobId = cellTwoSnapshot.get(random.nextInt(cellTwoSnapshot.size())); assertThat(cellOneService.containsCapacityUpdates(cellOneJobId)).isFalse(); assertThat(cellTwoService.containsCapacityUpdates(cellOneJobId)).isFalse(); JobCapacityUpdate cellOneUpdate = JobCapacityUpdate.newBuilder() .setJobId(cellOneJobId) .setCapacity(Capacity.newBuilder() .setMax(1).setDesired(2).setMax(3) .build()) .build(); AssertableSubscriber<Void> testSubscriber = service.updateJobCapacity(cellOneUpdate, UNDEFINED_CALL_METADATA).test(); testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS); testSubscriber.assertNoErrors(); testSubscriber.assertNoValues(); testSubscriber.assertCompleted(); assertThat(cellOneService.containsCapacityUpdates(cellOneJobId)).isTrue(); assertThat(cellTwoService.containsCapacityUpdates(cellOneJobId)).isFalse(); testSubscriber.unsubscribe(); JobCapacityUpdate cellTwoUpdate = JobCapacityUpdate.newBuilder() .setJobId(cellTwoJobId) .setCapacity(Capacity.newBuilder() .setMax(2).setDesired(2).setMax(2) .build()) .build(); testSubscriber = service.updateJobCapacity(cellTwoUpdate, UNDEFINED_CALL_METADATA).test(); testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS); testSubscriber.assertNoErrors(); testSubscriber.assertNoValues(); testSubscriber.assertCompleted(); assertThat(cellOneService.containsCapacityUpdates(cellTwoJobId)).isFalse(); assertThat(cellTwoService.containsCapacityUpdates(cellTwoJobId)).isTrue(); testSubscriber.unsubscribe(); }