Java Code Examples for rx.observers.AssertableSubscriber#awaitTerminalEvent()

@Test
public void grpcErrors() {
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(ok("one")));
    cellTwo.getServiceRegistry().addService(new HealthGrpc.HealthImplBase() {
        @Override
        public void check(HealthCheckRequest request, StreamObserver<HealthCheckResponse> responseObserver) {
            responseObserver.onError(Status.DEADLINE_EXCEEDED.asRuntimeException());
        }
    });

    AssertableSubscriber<HealthCheckResponse> subscriber = service.check(HealthCheckRequest.newBuilder().build()).test();
    subscriber.awaitTerminalEvent(10, TimeUnit.SECONDS);
    subscriber.assertNoErrors();
    subscriber.assertValueCount(1);
    HealthCheckResponse response = subscriber.getOnNextEvents().get(0);
    assertThat(response.getStatus()).isEqualTo(NOT_SERVING);
    assertThat(response.getDetailsCount()).isEqualTo(2);
    List<ServerStatus> errors = response.getDetailsList().stream()
            .filter(ServerStatus::hasError)
            .collect(Collectors.toList());
    assertThat(errors).hasSize(1);
    assertThat(errors.get(0).getError().getCell()).isEqualTo("two");
    assertThat(errors.get(0).getError().getErrorCode()).isEqualTo(DEADLINE_EXCEEDED.toString());
}
 

@Test
public void observeJobsStopsWhenAnyClientsTerminate() {
    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize()));

    final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test();

    final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-100").setStatus(ACCEPTED_STATE).build());
    final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-200").setStatus(ACCEPTED_STATE).build());
    cellOneUpdates.onNext(cellOneUpdate);
    cellTwoUpdates.onNext(cellTwoUpdate);

    testSubscriber.awaitValueCount(2, 1, TimeUnit.SECONDS);
    assertThat(testSubscriber.getOnErrorEvents()).isEmpty();
    assertThat(testSubscriber.isUnsubscribed()).isFalse();
    assertThat(testSubscriber.getCompletions()).isEqualTo(0);

    // a client finishes
    cellTwoUpdates.onCompleted();

    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    assertThat(testSubscriber.getOnErrorEvents()).isEmpty();
    assertThat(testSubscriber.isUnsubscribed()).isTrue();
    assertThat(testSubscriber.getCompletions()).isEqualTo(1);
}
 

@Test
public void getLoadBalancersWithOneFailingCell() {
    JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1);
    JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_1, LB_2);
    final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(
            Arrays.asList(jobLoadBalancer1, jobLoadBalancer2));

    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellTwo.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.INTERNAL));

    final AssertableSubscriber<GetAllLoadBalancersResult> resultSubscriber = service.getAllLoadBalancers(
            GetAllLoadBalancersRequest.newBuilder().setPage(Page.newBuilder().setPageSize(10)).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoValues();
    final List<Throwable> onErrorEvents = resultSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents).hasSize(1);
    assertThat(Status.fromThrowable(onErrorEvents.get(0))).isEqualTo(Status.INTERNAL);
}
 

@Test
public void addLoadBalancer() {
    JobLoadBalancer jobLoadBalancer1 = new JobLoadBalancer(JOB_1, LB_1);
    JobLoadBalancer jobLoadBalancer2 = new JobLoadBalancer(JOB_2, LB_2);
    final CellWithLoadBalancers cellWithLoadBalancersOne = new CellWithLoadBalancers(singletonList(jobLoadBalancer1));
    final CellWithLoadBalancers cellWithLoadBalancersTwo = new CellWithLoadBalancers(new ArrayList<>(singletonList(jobLoadBalancer2)));

    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.addLoadBalancer(
            AddLoadBalancerRequest.newBuilder().setJobId(JOB_2).setLoadBalancerId(LoadBalancerId.newBuilder().setId(LB_3).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.assertNoErrors();
    final List<GetJobLoadBalancersResult> onNextEvents = jobResults.getOnNextEvents();
    assertThat(onNextEvents.size()).isEqualTo(1);
    final List<LoadBalancerId> loadBalancersList = onNextEvents.get(0).getLoadBalancersList();
    assertThat(loadBalancersList.size()).isEqualTo(2);

    final List<String> resultLoadBalancerIds = loadBalancersList.stream().map(LoadBalancerId::getId).collect(Collectors.toList());
    assertThat(resultLoadBalancerIds).contains(LB_2, LB_3);
}
 

@Test
public void singleCell() {
    reset(connector);
    when(connector.getChannels()).thenReturn(Collections.singletonMap(
            new Cell("one", "1"), cellOne.getChannel()
    ));
    when(connector.getChannelForCell(any())).thenReturn(Optional.of(cellOne.getChannel()));

    HealthCheckResponse one = HealthCheckResponse.newBuilder()
            .setStatus(NOT_SERVING)
            .addDetails(ServerStatus.newBuilder()
                    .setDetails(Details.newBuilder()
                            .setCell("one")
                            .setLeader(false)
                            .setActive(true)
                            .setUptime(Durations.fromMillis(clock.wallTime()))
                            .build())
                    .build())
            .build();
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(one));

    AssertableSubscriber<HealthCheckResponse> subscriber = service.check(HealthCheckRequest.newBuilder().build()).test();
    subscriber.awaitTerminalEvent(10, TimeUnit.SECONDS);
    subscriber.assertNoErrors();
    subscriber.assertValueCount(1);
    HealthCheckResponse response = subscriber.getOnNextEvents().get(0);
    assertThat(response.getStatus()).isEqualTo(NOT_SERVING);
    assertThat(response.getDetailsCount()).isEqualTo(1);
    assertThat(response.getDetails(0).hasDetails()).isTrue();
    assertThat(response.getDetails(0).getDetails().getCell()).isEqualTo("one");
    assertThat(response.getDetails(0).getDetails().getLeader()).isFalse();
}
 

@Test
public void getJobLoadBalancersWithTwoFailingCell() {
    cellOne.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.NOT_FOUND));
    cellTwo.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.INTERNAL));

    final AssertableSubscriber<GetJobLoadBalancersResult> resultSubscriber = service.getLoadBalancers(
            JobId.newBuilder().setId(JOB_1).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoValues();
    final List<Throwable> onErrorEvents = resultSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents).hasSize(1);
    assertThat(Status.fromThrowable(onErrorEvents.get(0))).isEqualTo(Status.INTERNAL);
}
 

@Test
public void killJob() {
    Random random = new Random();
    List<Job> cellOneSnapshot = new ArrayList<>(dataGenerator.newServiceJobs(12, GrpcJobManagementModelConverters::toGrpcJob));
    List<Job> cellTwoSnapshot = new ArrayList<>(dataGenerator.newBatchJobs(7, GrpcJobManagementModelConverters::toGrpcJob));
    CellWithFixedJobsService cellOneService = new CellWithFixedJobsService(cellOneSnapshot, cellOneUpdates.serialize());
    CellWithFixedJobsService cellTwoService = new CellWithFixedJobsService(cellTwoSnapshot, cellTwoUpdates.serialize());
    cellOne.getServiceRegistry().addService(cellOneService);
    cellTwo.getServiceRegistry().addService(cellTwoService);

    Job killInCellOne = cellOneSnapshot.get(random.nextInt(cellOneSnapshot.size()));
    Job killInCellTwo = cellTwoSnapshot.get(random.nextInt(cellTwoSnapshot.size()));
    assertThat(cellOneService.currentJobs()).containsKey(killInCellOne.getId());
    assertThat(cellTwoService.currentJobs()).containsKey(killInCellTwo.getId());

    AssertableSubscriber<Void> testSubscriber = service.killJob(killInCellOne.getId(), UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    testSubscriber.assertNoValues();
    testSubscriber.assertCompleted();
    assertThat(cellOneService.currentJobs()).doesNotContainKey(killInCellOne.getId());
    assertThat(cellTwoService.currentJobs()).doesNotContainKey(killInCellOne.getId());
    testSubscriber.unsubscribe();

    testSubscriber = service.killJob(killInCellTwo.getId(), UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    testSubscriber.assertNoValues();
    testSubscriber.assertCompleted();
    assertThat(cellOneService.currentJobs()).doesNotContainKey(killInCellTwo.getId());
    assertThat(cellTwoService.currentJobs()).doesNotContainKey(killInCellTwo.getId());
    testSubscriber.unsubscribe();
}
 

@Test
public void findJobErrors() {
    cellOne.getServiceRegistry().addService(new CellWithFailingJobManagementService(NOT_FOUND));
    cellTwo.getServiceRegistry().addService(new CellWithFailingJobManagementService(UNAVAILABLE));

    AssertableSubscriber<Job> testSubscriber = service.findJob("any", UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    // transient errors have higher precedence than not found
    testSubscriber.assertError(StatusRuntimeException.class);
    assertThat(Status.fromThrowable(testSubscriber.getOnErrorEvents().get(0))).isEqualTo(UNAVAILABLE);
}
 

@Test
public void getJobLoadBalancersInvalidJobId() {
    cellOne.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.NOT_FOUND));
    cellTwo.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.NOT_FOUND));

    final AssertableSubscriber<GetJobLoadBalancersResult> resultSubscriber = service.getLoadBalancers(
            JobId.newBuilder().setId(JOB_1).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoValues();
    final List<Throwable> onErrorEvents = resultSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents).hasSize(1);
    assertThat(Status.fromThrowable(onErrorEvents.get(0))).isEqualTo(Status.NOT_FOUND);
}
 

@Test
public void observeJob() {
    String cellOneJobId = UUID.randomUUID().toString();
    String cellTwoJobId = UUID.randomUUID().toString();
    cellOne.getServiceRegistry().addService(new CellWithJobStream(cellOneJobId, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithJobStream(cellTwoJobId, cellTwoUpdates.serialize()));

    AssertableSubscriber<JobChangeNotification> subscriber1 = service.observeJob(cellOneJobId, UNDEFINED_CALL_METADATA).test();
    AssertableSubscriber<JobChangeNotification> subscriber2 = service.observeJob(cellTwoJobId, UNDEFINED_CALL_METADATA).test();

    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(ACCEPTED_STATE).build()));
    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(KILL_INITIATED_STATE).build()));
    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(FINISHED_STATE).build()));
    cellOneUpdates.onNext(toNotification(Job.newBuilder().setId(cellOneJobId).setStatus(ACCEPTED_STATE).build()));

    subscriber1.awaitValueCount(3, 5, TimeUnit.SECONDS);
    subscriber1.assertNoErrors();
    subscriber1.assertNotCompleted();
    assertThat(subscriber1.isUnsubscribed()).isFalse();

    subscriber2.assertNoErrors();
    subscriber2.assertNoValues();
    subscriber2.assertNotCompleted();

    cellTwoUpdates.onNext(toNotification(Task.newBuilder()
            .setId(cellTwoJobId + "-task1").setJobId(cellTwoJobId)
            .build())
    );
    subscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);
    subscriber2.assertNoErrors();
    subscriber2.assertNotCompleted();

    cellOneUpdates.onCompleted();

    subscriber1.awaitTerminalEvent(1, TimeUnit.SECONDS);
    assertThat(subscriber1.getOnErrorEvents()).isEmpty();
    assertThat(subscriber1.isUnsubscribed()).isTrue();
    assertThat(subscriber1.getCompletions()).isEqualTo(1);
}
 

@Test
public void consumerTransformerShouldPropagateCompletion() throws Exception {
  AssertableSubscriber<Object> subscriber =
      upstream.compose(Transformers.fromConsumer(consumer, scheduler)).test();

  upstream.onNext("hi");
  upstream.onCompleted();

  scheduler.triggerActions();

  subscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
  subscriber.assertCompleted();
}
 

@Test
public void getLoadBalancersWithTwoFailingCell() {
    cellOne.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.UNAVAILABLE));
    cellTwo.getServiceRegistry().addService(new CellWithFailingLoadBalancers(Status.INTERNAL));

    final AssertableSubscriber<GetAllLoadBalancersResult> resultSubscriber = service.getAllLoadBalancers(
            GetAllLoadBalancersRequest.newBuilder().setPage(Page.newBuilder().setPageSize(10)).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    resultSubscriber.assertNoValues();
    final List<Throwable> onErrorEvents = resultSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents).hasSize(1);
    assertThat(Status.fromThrowable(onErrorEvents.get(0))).isEqualTo(Status.INTERNAL);
}
 

@Test
public void findJobsAddsStackName() {
    Random random = new Random();
    final List<Job> cellSnapshot = new ArrayList<>();
    for (int i = 0; i < 5; i++) {
        cellSnapshot.addAll(dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob));
        cellSnapshot.addAll(dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob));
        clock.advanceTime(1, TimeUnit.MINUTES);
    }
    cell.getServiceRegistry().addService(new CellWithFixedJobsService(cellSnapshot, PublishSubject.create()));

    JobQuery query = JobQuery.newBuilder()
            .setPage(toGrpcPage(Page.unlimited()))
            .build();

    final AssertableSubscriber<JobQueryResult> testSubscriber = service.findJobs(query, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors().assertCompleted();
    testSubscriber.assertValueCount(1);
    final List<JobQueryResult> results = testSubscriber.getOnNextEvents();
    assertThat(results).hasSize(1);

    // expect stackName to have changed
    List<Job> expected = cellSnapshot.stream()
            .sorted(JobManagerCursors.jobCursorOrderComparator())
            .map(this::withStackName)
            .collect(Collectors.toList());
    assertThat(results.get(0).getItemsList()).containsExactlyElementsOf(expected);
}
 

@Test
public void shouldPropagateErrorsFromUpstream() throws Exception {
  final Throwable expected = new RuntimeException("expected");

  AssertableSubscriber<Integer> subscriber =
      input.compose(RxConnectables.toTransformer(connectable)).test();

  input.onError(expected);

  subscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
  subscriber.assertError(expected);
}
 

@Test
public void shouldPropagateErrorsFromConnectable() throws Exception {
  AssertableSubscriber<Integer> subscriber =
      input.compose(RxConnectables.toTransformer(connectable)).test();

  input.onNext("crash");

  subscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
  subscriber.assertFailureAndMessage(RuntimeException.class, "crashing!");
}
 

@Test
public void shouldPropagateCompletion() throws Exception {
  AssertableSubscriber<Integer> subscriber =
      input.compose(RxConnectables.toTransformer(connectable)).test();

  input.onNext("hi");
  input.onCompleted();

  subscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
  subscriber.assertCompleted();
}
 

@Test
public void shouldPropagateIncomingErrorsAsUnrecoverable() throws Exception {
  RxMobiusLoop<Integer, String, Boolean> loop =
      new RxMobiusLoop<>(builder, "", Collections.emptySet());
  PublishSubject<Integer> input = PublishSubject.create();

  AssertableSubscriber<String> subscriber = input.compose(loop).test();

  Exception expected = new RuntimeException("expected");
  input.onError(expected);
  subscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
  subscriber.assertError(new UnrecoverableIncomingException(expected));
  assertEquals(0, connection.valueCount());
}
 

@Test
public void findJobsMergesAllCellsIntoSingleResult() {
    Random random = new Random();
    final List<Job> cellOneSnapshot = new ArrayList<>();
    final List<Job> cellTwoSnapshot = new ArrayList<>();
    for (int i = 0; i < 5; i++) {
        cellOneSnapshot.addAll(
                dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        cellTwoSnapshot.addAll(
                dataGenerator.newBatchJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        cellOneSnapshot.addAll(
                dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        cellTwoSnapshot.addAll(
                dataGenerator.newServiceJobs(random.nextInt(10), GrpcJobManagementModelConverters::toGrpcJob)
        );
        clock.advanceTime(1, TimeUnit.MINUTES);
    }

    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(cellOneSnapshot, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(cellTwoSnapshot, cellTwoUpdates.serialize()));

    JobQuery query = JobQuery.newBuilder()
            .setPage(toGrpcPage(Page.unlimited()))
            .build();

    final AssertableSubscriber<JobQueryResult> testSubscriber = service.findJobs(query, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors().assertCompleted();
    testSubscriber.assertValueCount(1);
    final List<JobQueryResult> results = testSubscriber.getOnNextEvents();
    assertThat(results).hasSize(1);

    // expect stackName to have changed
    List<Job> expected = Stream.concat(cellOneSnapshot.stream(), cellTwoSnapshot.stream())
            .sorted(JobManagerCursors.jobCursorOrderComparator())
            .map(this::withStackName)
            .collect(Collectors.toList());
    assertThat(results.get(0).getItemsList()).containsExactlyElementsOf(expected);
}
 

@Test
public void deleteScalingPolicyIsIdempotent() {
    String jobId = UUID.randomUUID().toString();
    String policyId = UUID.randomUUID().toString();

    AWSApplicationAutoScalingAsync clientAsync = mock(AWSApplicationAutoScalingAsync.class);
    AWSAppScalingConfig config = mock(AWSAppScalingConfig.class);
    AWSAppAutoScalingClient autoScalingClient = new AWSAppAutoScalingClient(clientAsync, config, new NoopRegistry());

    // delete happens successfully on the first attempt
    AtomicBoolean isDeleted = new AtomicBoolean(false);
    when(clientAsync.deleteScalingPolicyAsync(any(), any())).thenAnswer(invocation -> {
        DeleteScalingPolicyRequest request = invocation.getArgument(0);
        AsyncHandler<DeleteScalingPolicyRequest, DeleteScalingPolicyResult> handler = invocation.getArgument(1);

        if (isDeleted.get()) {
            ObjectNotFoundException notFoundException = new ObjectNotFoundException(policyId + " does not exist");
            handler.onError(notFoundException);
            return Future.failed(notFoundException);
        }

        DeleteScalingPolicyResult resultSuccess = new DeleteScalingPolicyResult();
        HttpResponse successResponse = new HttpResponse(null, null);
        successResponse.setStatusCode(200);
        resultSuccess.setSdkHttpMetadata(SdkHttpMetadata.from(successResponse));

        isDeleted.set(true);
        handler.onSuccess(request, resultSuccess);
        return Future.successful(resultSuccess);
    });

    AssertableSubscriber<Void> firstCall = autoScalingClient.deleteScalingPolicy(policyId, jobId).test();
    firstCall.awaitTerminalEvent(2, TimeUnit.SECONDS);
    firstCall.assertNoErrors();
    firstCall.assertCompleted();
    verify(clientAsync, times(1)).deleteScalingPolicyAsync(any(), any());

    // second should complete fast when NotFound and not retry with exponential backoff
    AssertableSubscriber<Void> secondCall = autoScalingClient.deleteScalingPolicy(policyId, jobId).test();
    secondCall.awaitTerminalEvent(2, TimeUnit.SECONDS);
    secondCall.assertNoErrors();
    secondCall.assertCompleted();
    verify(clientAsync, times(2)).deleteScalingPolicyAsync(any(), any());
}
 

@Test
public void killTask() {
    Random random = new Random();
    List<Task> cellOneSnapshot = new ArrayList<>(dataGenerator.newServiceJobWithTasks());
    List<Task> cellTwoSnapshot = new ArrayList<>(dataGenerator.newBatchJobWithTasks());
    CellWithFixedTasksService cellOneService = new CellWithFixedTasksService(cellOneSnapshot);
    CellWithFixedTasksService cellTwoService = new CellWithFixedTasksService(cellTwoSnapshot);
    cellOne.getServiceRegistry().addService(cellOneService);
    cellTwo.getServiceRegistry().addService(cellTwoService);

    Task killInCellOne = cellOneSnapshot.get(random.nextInt(cellOneSnapshot.size()));
    Task killInCellTwo = cellTwoSnapshot.get(random.nextInt(cellTwoSnapshot.size()));
    assertThat(cellOneService.currentTasks()).containsKey(killInCellOne.getId());
    assertThat(cellTwoService.currentTasks()).containsKey(killInCellTwo.getId());

    TaskKillRequest cellOneRequest = TaskKillRequest.newBuilder()
            .setTaskId(killInCellOne.getId())
            .setShrink(false)
            .build();
    AssertableSubscriber<Void> testSubscriber = service.killTask(cellOneRequest, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    testSubscriber.assertNoValues();
    testSubscriber.assertCompleted();
    assertThat(cellOneService.currentTasks()).doesNotContainKey(killInCellOne.getId());
    assertThat(cellTwoService.currentTasks()).doesNotContainKey(killInCellOne.getId());
    testSubscriber.unsubscribe();

    TaskKillRequest cellTwoRequest = TaskKillRequest.newBuilder()
            .setTaskId(killInCellTwo.getId())
            .setShrink(false)
            .build();
    testSubscriber = service.killTask(cellTwoRequest, UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    testSubscriber.assertNoValues();
    testSubscriber.assertCompleted();
    assertThat(cellOneService.currentTasks()).doesNotContainKey(killInCellTwo.getId());
    assertThat(cellTwoService.currentTasks()).doesNotContainKey(killInCellTwo.getId());
    testSubscriber.unsubscribe();
}