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

@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 observeJobsErrorsWhenAnyClientsError() {
    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 emits an error
    cellTwoUpdates.onError(new RuntimeException("unexpected error"));

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

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

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

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

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

    final AssertableSubscriber<GetPolicyResult> testSubscriber = service.getAllScalingPolicies(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(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 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 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 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 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 getPolicyByIdFromTwoCells() {
    ScalingPolicyID policy1 = ScalingPolicyID.newBuilder().setId(POLICY_1).build();
    ScalingPolicyID policy2 = ScalingPolicyID.newBuilder().setId(POLICY_2).build();

    ScalingPolicyResult policyOneResult = ScalingPolicyResult.newBuilder().setId(policy1).build();
    ScalingPolicyResult policyTwoResult = ScalingPolicyResult.newBuilder().setId(policy2).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.getScalingPolicy(
            ScalingPolicyID.newBuilder().setId(POLICY_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);

    // Bad id. The current behavior is "INTERNAL: Completed without a response", but it will change to NOT_FOUND someday
    testSubscriber = service.getScalingPolicy(ScalingPolicyID.newBuilder().setId("badID").build(), JUNIT_REST_CALL_METADATA).test();
    testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS);
    testSubscriber.assertError(StatusRuntimeException.class);
    testSubscriber.assertNoValues();
    List<Throwable> onErrorEvents = testSubscriber.getOnErrorEvents();
    assertThat(onErrorEvents).isNotNull();
    assertThat(onErrorEvents).hasSize(1);
    assertThat(Status.fromThrowable(onErrorEvents.get(0)).getCode()).isEqualTo(Status.CANCELLED.getCode());
}
 

@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 observeJobsWaitsForAllMarkers() {
    final List<Job> cellOneSnapshot = Arrays.asList(
            Job.newBuilder().setId("cell-1-job-1").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-1-job-2").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-1-job-3").setStatus(ACCEPTED_STATE).build()
    );
    final List<Job> cellTwoSnapshot = Arrays.asList(
            Job.newBuilder().setId("cell-2-job-1").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-2-job-2").setStatus(ACCEPTED_STATE).build(),
            Job.newBuilder().setId("cell-2-job-3").setStatus(ACCEPTED_STATE).build()
    );

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

    final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test();
    List<JobChangeNotification> expected = Stream.concat(
            cellOneSnapshot.stream().map(this::toNotification).map(this::withStackName),
            cellTwoSnapshot.stream().map(this::toNotification).map(this::withStackName)
    ).collect(Collectors.toList());
    // single marker for all cells
    final JobChangeNotification mergedMarker = JobChangeNotification.newBuilder().setSnapshotEnd(SnapshotEnd.newBuilder()).build();
    expected.add(mergedMarker);

    testSubscriber.awaitValueCount(7, 1, TimeUnit.SECONDS);
    List<JobChangeNotification> onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).last().isEqualTo(mergedMarker);
    assertThat(onNextEvents).containsExactlyInAnyOrder(expected.toArray(new JobChangeNotification[expected.size()]));

    // more updates are flowing
    final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-10").setStatus(ACCEPTED_STATE).build());
    final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-10").setStatus(ACCEPTED_STATE).build());
    cellOneUpdates.onNext(cellOneUpdate);
    cellTwoUpdates.onNext(cellTwoUpdate);

    testSubscriber.awaitValueCount(9, 1, TimeUnit.SECONDS);
    onNextEvents = testSubscriber.getOnNextEvents();
    assertThat(onNextEvents).last().isNotEqualTo(mergedMarker);
    assertThat(onNextEvents).contains(withStackName(cellOneUpdate), withStackName(cellTwoUpdate));
}
 

@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 createAndDeletePolicyForJobIdFromTwoCells() {
    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);
    cellTwo.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);
    assertThat(testSubscriber.getOnErrorEvents().isEmpty()).isTrue();

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

    String newPolicyId = onNextEvents.get(0).getId();
    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);

    DeletePolicyRequest deletePolicyRequest = DeletePolicyRequest.newBuilder().setId(ScalingPolicyID.newBuilder().setId(newPolicyId).build()).build();
    AssertableSubscriber<Void> testSubscriber3 = service.deleteAutoScalingPolicy(deletePolicyRequest, JUNIT_REST_CALL_METADATA).test();
    testSubscriber3.assertNoErrors();

    AssertableSubscriber<GetPolicyResult> testSubscriber4 = service.getJobScalingPolicies(JobId.newBuilder().setId(JOB_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);
}