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

@Test
public void getJobLoadBalancersWithOneFailingCell() {
    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<GetJobLoadBalancersResult> resultSubscriber = service.getLoadBalancers(
            JobId.newBuilder().setId(JOB_1).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    final List<GetJobLoadBalancersResult> onNextEvents = resultSubscriber.getOnNextEvents();
    assertThat(onNextEvents).hasSize(1);
    final List<LoadBalancerId> loadBalancersList = onNextEvents.get(0).getLoadBalancersList();
    final List<String> resultLoadBalancers = loadBalancersList.stream().map(loadBalancerId -> loadBalancerId.getId()).collect(Collectors.toList());
    assertThat(resultLoadBalancers).contains(LB_1, LB_2);
}
 

@Test
public void getLoadBalancersForJob() {
    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(singletonList(jobLoadBalancer2));
    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellTwo.getServiceRegistry().addService(cellWithLoadBalancersTwo);

    final AssertableSubscriber<GetJobLoadBalancersResult> resultSubscriber = service.getLoadBalancers(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    resultSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    final List<GetJobLoadBalancersResult> onNextEvents = resultSubscriber.getOnNextEvents();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents.get(0).getLoadBalancersCount()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getLoadBalancers(0).getId()).isEqualTo(LB_2);
}
 

@Test
public void oneCellFailing() {
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(ok("one")));
    cellTwo.getServiceRegistry().addService(new CellWithHealthStatus(failing("two")));

    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);
    assertThat(response.getDetails(0).hasDetails()).isTrue();
    assertThat(response.getDetails(1).hasDetails()).isTrue();
    Set<String> cellsSeen = response.getDetailsList().stream()
            .map(s -> s.getDetails().getCell())
            .collect(Collectors.toSet());
    assertThat(cellsSeen).contains("one", "two");
}
 

@Test
public void allCellsOK() {
    cellOne.getServiceRegistry().addService(new CellWithHealthStatus(ok("one")));
    cellTwo.getServiceRegistry().addService(new CellWithHealthStatus(ok("two")));

    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(SERVING);
    assertThat(response.getDetailsCount()).isEqualTo(2);
    assertThat(response.getDetails(0).hasDetails()).isTrue();
    assertThat(response.getDetails(1).hasDetails()).isTrue();
    Set<String> cellsSeen = response.getDetailsList().stream()
            .map(s -> s.getDetails().getCell())
            .collect(Collectors.toSet());
    assertThat(cellsSeen).contains("one", "two");
}
 

@Test
public void onErrorIsNotSentAfterOnCompleted() {
    final RateLimitedBatcher<BatchableOperationMock, String> batcher = buildBatcher(minimumTimeInQueueMs);
    final Subject<BatchableOperationMock, BatchableOperationMock> updates = PublishSubject.<BatchableOperationMock>create().toSerialized();

    final AssertableSubscriber<Batch<BatchableOperationMock, String>> subscriber = updates.lift(batcher).test();
    testScheduler.triggerActions();
    subscriber.assertNoTerminalEvent().assertNoValues();

    updates.onCompleted();
    // onCompleted is sent after pending items are drained
    testScheduler.advanceTimeBy(minimumTimeInQueueMs, TimeUnit.MILLISECONDS);
    subscriber.assertNoErrors().assertCompleted();

    updates.onError(new RuntimeException("some problem"));
    subscriber.assertNoErrors();
}
 

@Test
public void getAllLoadBalancersNoPagination() {
    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(singletonList(jobLoadBalancer2));
    cellOne.getServiceRegistry().addService(cellWithLoadBalancersOne);
    cellTwo.getServiceRegistry().addService(cellWithLoadBalancersTwo);

    final AssertableSubscriber<GetAllLoadBalancersResult> resultSubscriber = service.getAllLoadBalancers(
            GetAllLoadBalancersRequest.newBuilder().setPage(Page.newBuilder().setPageNumber(0).setPageSize(10)).build(),
            JUNIT_REST_CALL_METADATA
    ).test();
    resultSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    resultSubscriber.assertNoErrors();
    final List<GetAllLoadBalancersResult> onNextEvents = resultSubscriber.getOnNextEvents();
    assertThat(onNextEvents).hasSize(1);
    assertThat(onNextEvents.get(0).getJobLoadBalancersCount()).isEqualTo(2);
}
 

@Test
public void findTask() {
    Random random = new Random();
    List<Task> cellOneSnapshot = new ArrayList<>(dataGenerator.newServiceJobWithTasks());
    cellOne.getServiceRegistry().addService(new CellWithFixedTasksService(cellOneSnapshot));
    cellTwo.getServiceRegistry().addService(new CellWithFixedTasksService(Collections.emptyList()));

    Task expected = withStackName(cellOneSnapshot.get(random.nextInt(cellOneSnapshot.size())));
    AssertableSubscriber<Task> testSubscriber = service.findTask(expected.getId(), UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    testSubscriber.assertValueCount(1);
    testSubscriber.assertValue(expected);
}
 

@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);
}
 

@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 startModelAndEffects() throws Exception {
  RxMobiusLoop<Integer, String, Boolean> loop =
      new RxMobiusLoop<>(builder, "StartModel", ImmutableSet.of(true, false));
  final AssertableSubscriber<String> subscriber = Observable.just(1).compose(loop).test();

  waitForSubscriberValueCount(subscriber, 2);
  subscriber.assertValues("StartModel", "StartModel1");
  subscriber.assertNoErrors();
  assertEquals(2, connection.valueCount());
  connection.assertValuesInAnyOrder(true, false);
}
 

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

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

    // TODO: make it easier to extract the Deadline for each cell call
    Thread.sleep(2 * GRPC_REQUEST_TIMEOUT_MS);

    subscriber.assertNoErrors();
    subscriber.assertNotCompleted();
    assertThat(subscriber.isUnsubscribed()).isFalse();
}
 

@Test
public void getPoliciesForJobFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).setJobId(JOB_1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).setJobId(JOB_2).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithPolicies cellTwoService = new CellWithPolicies(Collections.singletonList(policyTwoResult));

    cellOne.getServiceRegistry().addService(cellOneService);
    cellTwo.getServiceRegistry().addService(cellTwoService);

    AssertableSubscriber<GetPolicyResult> testSubscriber = service.getJobScalingPolicies(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();

    testSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItemsCount()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);

    // Bad policy id, currently each Cell returns an empty result
    testSubscriber = service.getJobScalingPolicies(JobId.newBuilder().setId("badID").build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getItemsCount()).isEqualTo(0);
}
 

@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 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 findJob() {
    Random random = new Random();
    List<Job> cellOneSnapshot = new ArrayList<>(dataGenerator.newServiceJobs(10, GrpcJobManagementModelConverters::toGrpcJob));
    cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(cellOneSnapshot, cellOneUpdates.serialize()));
    cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize()));

    Job expected = withStackName(cellOneSnapshot.get(random.nextInt(cellOneSnapshot.size())));
    AssertableSubscriber<Job> testSubscriber = service.findJob(expected.getId(), UNDEFINED_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();
    testSubscriber.assertValueCount(1);
    testSubscriber.assertValue(expected);
}
 

@Test
public void reconciliationErrorsDontHaltOthers() {
    String jobId = UUID.randomUUID().toString();
    String loadBalancerId = "lb-" + UUID.randomUUID().toString();
    ThreadLocalRandom random = ThreadLocalRandom.current();
    int batchSize = random.nextInt(3, 10);

    defaultStubs();
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, jobId);
    List<Task> tasks = LoadBalancerTests.buildTasksStarted(batchSize, jobId);
    when(v3JobOperations.getTasks(jobId)).thenReturn(tasks);

    LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS);
    DefaultLoadBalancerService service = new DefaultLoadBalancerService(
            runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler);

    AssertableSubscriber<Batch<TargetStateBatchable, String>> testSubscriber = service.events().test();

    assertTrue(service.addLoadBalancer(jobId, loadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(jobId).toBlocking().first()).isEqualTo(loadBalancerId);

    reconcilerEvents.onError(new RuntimeException("first error"));
    testScheduler.triggerActions();
    testSubscriber.assertNoErrors();

    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
    reconcilerEvents.onError(new RuntimeException("second error"));

    testSubscriber.assertNoErrors().assertValueCount(1);
    verify(client).registerAll(eq(loadBalancerId), argThat(targets -> targets != null && targets.size() == batchSize));
    verify(client, never()).deregisterAll(any(), any());
    verifyReconcilerIgnore(loadBalancerId, LoadBalancerTests.ipAddresses(tasks));
}
 

@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();
}
 

@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();
}
 

@Test
public void createAndUpdatePolicyForJobIdFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).setJobId(JOB_1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).setJobId(JOB_2).build();

    CellWithPolicies cellOneService = new CellWithPolicies(Collections.singletonList(policyOneResult));
    CellWithPolicies cellTwoService = new CellWithPolicies(Collections.singletonList(policyTwoResult));

    CellWithJobIds cellOneJobsService = new CellWithJobIds(Collections.singletonList(JOB_1));
    CellWithJobIds cellTwoJobsService = new CellWithJobIds(Collections.singletonList(JOB_2));
    cellOne.getServiceRegistry().addService(cellOneService);
    cellOne.getServiceRegistry().addService(cellOneJobsService);
    cellTwo.getServiceRegistry().addService(cellTwoService);
    cellTwo.getServiceRegistry().addService(cellTwoJobsService);

    AssertableSubscriber<ScalingPolicyID> testSubscriber = service.setAutoScalingPolicy(PutPolicyRequest.newBuilder().setJobId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitValueCount(1, 1, TimeUnit.SECONDS);
    testSubscriber.assertNoErrors();

    List<ScalingPolicyID> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).isNotNull();
    assertThat(onNextEvents.size()).isEqualTo(1);
    assertThat(onNextEvents.get(0).getId()).isNotEmpty();

    AssertableSubscriber<GetPolicyResult> testSubscriber2 = service.getJobScalingPolicies(JobId.newBuilder().setId(JOB_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents1 = testSubscriber2.getOnNextEvents();
    assertThat(onNextEvents1).isNotNull();
    assertThat(onNextEvents1.size()).isEqualTo(1);
    assertThat(onNextEvents1.get(0).getItemsCount()).isEqualTo(2);
    assertThat(onNextEvents1.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents1.get(0).getItems(1).getJobId()).isEqualTo(JOB_2);

    UpdatePolicyRequest updatePolicyRequest = UpdatePolicyRequest.newBuilder()
            .setPolicyId(ScalingPolicyID.newBuilder().setId(POLICY_2))
            .setScalingPolicy(ScalingPolicy.newBuilder()
                    .setTargetPolicyDescriptor(TargetTrackingPolicyDescriptor.newBuilder()
                            .setTargetValue(DoubleValue.newBuilder().setValue(100.0).build()).build()).build()).build();

    AssertableSubscriber<Void> testSubscriber3 = service.updateAutoScalingPolicy(updatePolicyRequest, JUNIT_REST_CALL_METADATA).test();
    testSubscriber3.assertNoErrors();

    AssertableSubscriber<GetPolicyResult> testSubscriber4 = service.getScalingPolicy(ScalingPolicyID.newBuilder().setId(POLICY_2).build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber2.awaitValueCount(1, 1, TimeUnit.SECONDS);

    List<GetPolicyResult> onNextEvents2 = testSubscriber4.getOnNextEvents();
    assertThat(onNextEvents2).isNotNull();
    assertThat(onNextEvents2.size()).isEqualTo(1);
    assertThat(onNextEvents2.get(0).getItemsCount()).isEqualTo(1);
    assertThat(onNextEvents2.get(0).getItems(0).getJobId()).isEqualTo(JOB_2);
    assertThat(onNextEvents2.get(0).getItems(0).getId().getId()).isEqualTo(POLICY_2);
    ScalingPolicy scalingPolicy = onNextEvents2.get(0).getItems(0).getScalingPolicy();
    double updatedValue = scalingPolicy.getTargetPolicyDescriptor().getTargetValue().getValue();
    assertThat(updatedValue).isEqualTo(100);
}
 

@Test
public void deleteScalableTargetIsIdempotent() {
    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());

    AtomicBoolean isDeleted = new AtomicBoolean(false);
    when(clientAsync.deregisterScalableTargetAsync(any(), any())).thenAnswer(invocation -> {
        DeregisterScalableTargetRequest request = invocation.getArgument(0);
        AsyncHandler<DeregisterScalableTargetRequest, DeregisterScalableTargetResult> handler = invocation.getArgument(1);
        if (isDeleted.get()) {
            ObjectNotFoundException notFoundException = new ObjectNotFoundException(policyId + " does not exist");
            handler.onError(notFoundException);
            return Future.failed(notFoundException);
        }

        DeregisterScalableTargetResult resultSuccess = new DeregisterScalableTargetResult();
        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.deleteScalableTarget(jobId).test();
    firstCall.awaitTerminalEvent(2, TimeUnit.SECONDS);
    firstCall.assertNoErrors();
    firstCall.assertCompleted();
    verify(clientAsync, times(1)).deregisterScalableTargetAsync(any(), any());

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