Java Code Examples for rx.subjects.PublishSubject#onNext()

/**
 * Provides a {@link TransformedProperty} that is "map" transformation of the given {@link Property}.
 * @param property the property
 * @param mapFunction the function that does the mapping for the transformation
 * @param <T> type of the property
 * @return the transformed property
 */
public static <T, U> TransformedProperty<U> map(Property<T> property, Func1<T, U> mapFunction) {
    Assert.requireNonNull(property, "property");
    Assert.requireNonNull(mapFunction, "mapFunction");

    final PublishSubject<T> reactiveObservable = PublishSubject.create();

    Subscription basicSubscription = property.onChanged(new ValueChangeListener<T>() {
        @Override
        public void valueChanged(ValueChangeEvent<? extends T> evt) {
            reactiveObservable.onNext(evt.getNewValue());
        }
    });

    TransformedPropertyImpl result = new TransformedPropertyImpl<>(basicSubscription);

    Observable<U> transformedObservable = reactiveObservable.map(mapFunction);
    transformedObservable.subscribe(result);

    reactiveObservable.onNext(property.get());

    return result;
}
 

public static <K, V> LegacyTcpPushServer<KeyValuePair<K, V>> infiniteStreamLegacyTcpNestedMantisGroup(ServerConfig<KeyValuePair<K, V>> config,
                                                                                                      Observable<Observable<MantisGroup<K, V>>> go,
                                                                                                      long groupExpirySeconds, final Func1<K, byte[]> keyEncoder,
                                                                                                      HashFunction hashFunction) {
    final PublishSubject<String> serverSignals = PublishSubject.create();
    final String serverName = config.getName();
    Action0 onComplete = new Action0() {
        @Override
        public void call() {
            serverSignals.onNext("ILLEGAL_STATE_COMPLETED");
            throw new IllegalStateException("OnComplete signal received, Server: " + serverName + " is pushing an infinite stream, should not complete");
        }
    };

    Action1<Throwable> onError = new Action1<Throwable>() {
        @Override
        public void call(Throwable t) {
            serverSignals.onError(t);
        }
    };

    PushTrigger<KeyValuePair<K, V>> trigger = ObservableTrigger.oomgo(serverName, go, onComplete, onError, groupExpirySeconds,
            keyEncoder, hashFunction);
    return new LegacyTcpPushServer<KeyValuePair<K, V>>(trigger, config, serverSignals);
}
 

/**
 * Provides a {@link TransformedProperty} that is "debounce" transformation of the given {@link Property}.
 * @param property the property
 * @param timeout timeout for the "debounce" transformation
 * @param unit time unit for the "debounce" transformation
 * @param <T> type of the property
 * @return the transformed property
 */
public static <T> TransformedProperty<T> debounce(Property<T> property, long timeout, TimeUnit unit) {
    Assert.requireNonNull(property, "property");
    Assert.requireNonNull(unit, "unit");

    final PublishSubject<T> reactiveObservable = PublishSubject.create();

    Subscription basicSubscription = property.onChanged(new ValueChangeListener<T>() {
        @Override
        public void valueChanged(ValueChangeEvent<? extends T> evt) {
            reactiveObservable.onNext(evt.getNewValue());
        }
    });

    TransformedPropertyImpl result = new TransformedPropertyImpl<>(basicSubscription);

    Observable<T> transformedObservable = reactiveObservable.debounce(timeout, unit);
    transformedObservable.subscribe(result);

    reactiveObservable.onNext(property.get());

    return result;
}
 

@Test
public void testSampleWindowIsConstantDuration() {
    @SuppressWarnings("unchecked")
    Observer<Integer> observer = mock(Observer.class);
    TestScheduler s = new TestScheduler();
    PublishSubject<Integer> o = PublishSubject.create();
    o.compose(Transformers.<Integer> sampleFirst(1000, TimeUnit.MILLISECONDS, s))
            .subscribe(observer);

    // send events with simulated time increments
    s.advanceTimeTo(0, TimeUnit.MILLISECONDS);
    o.onNext(1);
    s.advanceTimeTo(1200, TimeUnit.MILLISECONDS);
    o.onNext(2);
    s.advanceTimeTo(2100, TimeUnit.MILLISECONDS);
    o.onNext(3);
    o.onCompleted();

    InOrder inOrder = inOrder(observer);
    inOrder.verify(observer).onNext(1);
    inOrder.verify(observer).onNext(2);
    inOrder.verify(observer).onNext(3);
    inOrder.verify(observer).onCompleted();
    inOrder.verifyNoMoreInteractions();
}
 

@Test public void getAppsAfterRetryClickWithLoading() {
  //Given an initialized ListStoreAppsPresenter with a STORE_ID and a limit of apps
  //When onCreate lifecycle call event happens
  //and the user is on the no network view
  //and clicks on the retry button

  PublishSubject<Void> retryClickEvent = PublishSubject.create();
  when(view.getRetryEvent()).thenReturn(retryClickEvent);
  when(appCenter.loadNextApps(STORE_ID_TEST, LIMIT_APPS_TEST)).thenReturn(
      Single.just(appsModelLoading));
  listStoreAppsPresenter.present();
  lifecycleEvent.onNext(View.LifecycleEvent.CREATE);
  retryClickEvent.onNext(null);

  //then a loading should be shown in the UI
  verify(view).showStartingLoading();
  //and a request apps should be done
  verify(appCenter).loadNextApps(STORE_ID_TEST, LIMIT_APPS_TEST);
  //and apps model is loading
  //hide loading
  verify(view).hideLoading();
  //should do nothing else
  verify(view, never()).showNetworkError();
  verify(view, never()).showGenericError();
  verify(view, never()).addApps(appsModelLoading.getList());
}
 

@Test public void loadAppsAfterRefreshWithNetworkError() {
  //Given an initialized ListStoreAppsPresenter with a STORE_ID and a limit of apps
  //When onCreate lifecycle call event happens
  //and the view is refreshed by the user
  PublishSubject<Void> refreshEvent = PublishSubject.create();
  when(view.getRefreshEvent()).thenReturn(refreshEvent);

  when(appCenter.loadFreshApps(STORE_ID_TEST, LIMIT_APPS_TEST)).thenReturn(
      Single.just(appsModelWithNetworkError));

  listStoreAppsPresenter.present();
  lifecycleEvent.onNext(View.LifecycleEvent.CREATE);
  refreshEvent.onNext(null);

  //Then when new apps are requested
  verify(appCenter).loadFreshApps(STORE_ID_TEST, LIMIT_APPS_TEST);
  //and hide refresh loading
  verify(view).hideRefreshLoading();
  //and error - Generic error- show Generic error
  verify(view).showNetworkError();
}
 

@Test
public void testErrorsPassedThroughToOperator() {
    PublishSubject<Integer> subject = PublishSubject.create();
    Recorder recorder1 = new Recorder();
    Recorder recorder2 = new Recorder();
    subject.subscribe(recorder1);
    subject.subscribe(recorder2);
    subject.onNext(1);
    assertEquals(asList(1), recorder1.list());
    assertEquals(asList(1), recorder2.list());
    subject.onNext(2);
    assertEquals(asList(1, 2), recorder1.list());
    assertEquals(asList(1, 2), recorder2.list());
    Exception e = new Exception("boo");
    assertTrue(recorder1.errors().isEmpty());
    assertTrue(recorder2.errors().isEmpty());
    subject.onError(e);
    assertEquals(asList(e), recorder1.errors());
    assertEquals(asList(e), recorder2.errors());
}
 

public static void main(String[] args) throws Exception {
PublishSubject<Integer> ps = PublishSubject.create();

ps.publish(o -> 
    o.mergeWith(
        o.switchMap(e -> 
             Observable.just(1).delay(200, TimeUnit.MILLISECONDS)
            .ignoreElements()
            .doOnCompleted(() -> System.out.println("Timeout action: " + e))
        )
    )
).subscribe(System.out::println);

ps.onNext(1);
ps.onNext(2);

Thread.sleep(100);

ps.onNext(3);

Thread.sleep(250);

ps.onNext(4);

Thread.sleep(250);
}
 

@Test public void loadAppsAfterReachingBottomWithLoading() {
  //Given an initialized ListStoreAppsPresenter with a STORE_ID and a limit of apps
  //When onCreate lifecycle call event happens
  //and the view reaches the bottom
  PublishSubject<Object> reachedBottomEvent = PublishSubject.create();
  when(view.reachesBottom()).thenReturn(reachedBottomEvent);

  when(appCenter.loadNextApps(STORE_ID_TEST, LIMIT_APPS_TEST)).thenReturn(
      Single.just(appsModelLoading));

  listStoreAppsPresenter.present();
  lifecycleEvent.onNext(View.LifecycleEvent.CREATE);
  reachedBottomEvent.onNext(null);

  //then should show loading
  verify(view).showLoading();
  //and request the next apps
  verify(appCenter).loadNextApps(STORE_ID_TEST, LIMIT_APPS_TEST);
  //hide loading
  verify(view).hideLoading();
  //and should not show new apps in the UI or errors
  verify(view, never()).addApps(appsModelLoading.getList());
  verify(view, never()).showGenericError();
  verify(view, never()).showNetworkError();
}
 

@Test
public void should_fail_on_second_subscription() {
    // given
    PublishSubject<Integer> subject = PublishSubject.create();
    Observable<Integer> limitedObservable = SubscriptionLimiter.limitSubscriptions(1, subject);
    TestSubscriber<Integer> subscriber = new TestSubscriber<>();
    TestSubscriber<Integer> subscriber2 = new TestSubscriber<>();
    limitedObservable.subscribe(subscriber);
    // when
    limitedObservable.subscribe(subscriber2);
    subject.onNext(123);
    // then
    assertThat(subscriber2.getOnNextEvents()).isEmpty();
    assertThat(subscriber2.getOnErrorEvents()).hasSize(1);

}
 

@Test
public void errorsPropagateDownTheChain() {
  PublishSubject<Integer> upstream = PublishSubject.create();
  FlatMapCompletable<Integer, SomeOtherType> flatMapCompletable =
      FlatMapCompletable.createForAction(
          new Action1<Integer>() {
            @Override
            public void call(Integer integer) {
              throw new IllegalArgumentException();
            }
          });
  Observable<SomeOtherType> observable = upstream.compose(flatMapCompletable);

  TestSubscriber<SomeOtherType> subscriber = new TestSubscriber<>();
  observable.subscribe(subscriber);

  upstream.onNext(1);
  subscriber.assertNoValues();
  subscriber.assertError(IllegalArgumentException.class);
}
 

@Test
public void should_allow_one_subscription() {
    // given
    PublishSubject<Integer> subject = PublishSubject.create();
    Observable<Integer> limitedObservable = SubscriptionLimiter.limitSubscriptions(1, subject);
    TestSubscriber<Integer> subscriber = new TestSubscriber<>();
    // when
    limitedObservable.subscribe(subscriber);
    subject.onNext(123);
    // then
    assertThat(subscriber.getOnNextEvents()).hasSize(1).contains(123);
}
 

@Test
public void testSampleDoesNotRepeatEmissions() throws InterruptedException {
    PublishSubject<Integer> subject = PublishSubject.create();
    @SuppressWarnings("unchecked")
    final List<Integer> list = Mockito.mock(List.class);
    Subscription sub = subject
    // sample
            .sample(100, TimeUnit.MILLISECONDS)
            // subscribe and record emissions
            .subscribe(new Action1<Integer>() {
                @Override
                public void call(Integer n) {
                    list.add(n);
                }
            });
    // these items should be emitted before the first sample window closes
    subject.onNext(1);
    subject.onNext(2);
    subject.onNext(3);

    // wait for at least one more sample window to have passed
    Thread.sleep(250);

    // check that sample only reported the last item (3) once
    InOrder inOrder = Mockito.inOrder(list);
    inOrder.verify(list).add(3);
    inOrder.verifyNoMoreInteractions();

    sub.unsubscribe();
}
 

@MainThread public void onRequestPermissionsResult(@NonNull String[] permissions,
                                                   @NonNull int[] grantResults) {
    assertInitialized(rxPermissions);
    HashMap<String, PublishSubject<Boolean>> subjects = rxPermissions.subjects;
    for (int size = permissions.length, i = 0; i < size; i++) {
        PublishSubject<Boolean> subject = subjects.get(permissions[i]);
        assertInitialized(subject);
        subject.onNext(grantResults[i] == PackageManager.PERMISSION_GRANTED);
    }
    resultDelivered = true;
}
 

public synchronized void forceComplete(String key) {
    PublishSubject<Integer> takeUntil = takeUntilSubjects.get(key);
    if (takeUntil != null) {
        takeUntil.onNext(1); // complete observable
    } else {
        logger.debug("Nothing to force complete, key doesn't exist: " + key);
    }
}
 

public static <T, S> PushServerSse<T, S> infiniteStreamSse(ServerConfig<T> config, Observable<T> o,
                                                           Func2<Map<String, List<String>>, S, Void> requestPreprocessor,
                                                           Func2<Map<String, List<String>>, S, Void> requestPostprocessor,
                                                           final Func2<Map<String, List<String>>, S, Void> subscribeProcessor,
                                                           S state, boolean supportLegacyMetrics) {

    final String serverName = config.getName();
    final PublishSubject<String> serverSignals = PublishSubject.create();
    Action0 onComplete = new Action0() {
        @Override
        public void call() {
            serverSignals.onNext("ILLEGAL_STATE_COMPLETED");
            throw new IllegalStateException("OnComplete signal received, Server: " + serverName + " is pushing an infinite stream, should not complete");
        }
    };

    Action1<Throwable> onError = new Action1<Throwable>() {
        @Override
        public void call(Throwable t) {
            serverSignals.onError(t);
        }
    };

    PushTrigger<T> trigger = ObservableTrigger.o(serverName, o, onComplete, onError);

    return new PushServerSse<T, S>(trigger, config, serverSignals,
            requestPreprocessor, requestPostprocessor,
            subscribeProcessor, state, supportLegacyMetrics);
}
 

@Test
public void newTasksGetRegistered() {
    String jobId = UUID.randomUUID().toString();
    String taskId = UUID.randomUUID().toString();
    String loadBalancerId = "lb-" + UUID.randomUUID().toString();
    PublishSubject<JobManagerEvent<?>> taskEvents = PublishSubject.create();

    when(client.registerAll(any(), any())).thenReturn(Completable.complete());
    when(client.deregisterAll(any(), any())).thenReturn(Completable.complete());
    when(v3JobOperations.observeJobs()).thenReturn(taskEvents);
    when(v3JobOperations.getTasks(jobId)).thenReturn(Collections.emptyList());
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, jobId);

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

    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();

    Task launched = ServiceJobTask.newBuilder()
            .withJobId(jobId)
            .withId(taskId)
            .withStatus(TaskStatus.newBuilder().withState(TaskState.Launched).build())
            .build();

    Task startingWithIp = launched.toBuilder()
            .withStatus(TaskStatus.newBuilder().withState(TaskState.StartInitiated).build())
            .withTaskContext(CollectionsExt.asMap(
                    TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "1.2.3.4"
            )).build();

    Task started = startingWithIp.toBuilder()
            .withStatus(TaskStatus.newBuilder().withState(TaskState.Started).build())
            .build();
    LoadBalancerTargetState expectedTarget = new LoadBalancerTargetState(
            new LoadBalancerTarget(loadBalancerId, started.getId(), "1.2.3.4"),
            REGISTERED
    );

    // events with no state transition gets ignored
    taskEvents.onNext(TaskUpdateEvent.newTask(null, launched, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();

    // events to !Started states get ignored
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, startingWithIp, launched, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();

    // finally detect the task is UP and gets registered
    taskEvents.onNext(TaskUpdateEvent.taskChange(null, started, startingWithIp, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(1);
    assertThat(loadBalancerStore.getLoadBalancerTargets(loadBalancerId).collectList().block())
            .contains(expectedTarget);
    verify(client).registerAll(eq(loadBalancerId), argThat(set -> set.contains("1.2.3.4")));
    verify(client, never()).deregisterAll(eq(loadBalancerId), any());
    verifyReconcilerIgnore(loadBalancerId, "1.2.3.4");
}
 

@Test
public void movedTasks() {
    String taskId = UUID.randomUUID().toString();
    String sourceJobId = UUID.randomUUID().toString();
    String targetJobId = UUID.randomUUID().toString();
    String sourceLoadBalancerId = "lb-" + UUID.randomUUID().toString();
    String targetLoadBalancerId = "lb-" + UUID.randomUUID().toString();
    String commonLoadBalancerId = "lb-" + UUID.randomUUID().toString();
    PublishSubject<JobManagerEvent<?>> taskEvents = PublishSubject.create();

    when(client.registerAll(any(), any())).thenReturn(Completable.complete());
    when(client.deregisterAll(any(), any())).thenReturn(Completable.complete());
    when(v3JobOperations.observeJobs()).thenReturn(taskEvents);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, sourceJobId);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, targetJobId);

    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(sourceJobId, sourceLoadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertTrue(service.addLoadBalancer(sourceJobId, commonLoadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(sourceJobId).toBlocking().toIterable())
            .containsExactlyInAnyOrder(sourceLoadBalancerId, commonLoadBalancerId);

    assertTrue(service.addLoadBalancer(targetJobId, targetLoadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertTrue(service.addLoadBalancer(targetJobId, commonLoadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(targetJobId).toBlocking().toIterable())
            .containsExactlyInAnyOrder(targetLoadBalancerId, commonLoadBalancerId);

    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();

    Task moved = ServiceJobTask.newBuilder()
            .withJobId(targetJobId)
            .withId(taskId)
            .withStatus(TaskStatus.newBuilder().withState(TaskState.Started).build())
            .withTaskContext(CollectionsExt.asMap(
                    TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "1.2.3.4",
                    TaskAttributes.TASK_ATTRIBUTES_MOVED_FROM_JOB, sourceJobId
            )).build();

    // detect the task is moved, gets deregistered from the source and registered on the target
    taskEvents.onNext(TaskUpdateEvent.newTaskFromAnotherJob(null, moved, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(2);
    verify(client).registerAll(eq(targetLoadBalancerId), argThat(set -> set.contains("1.2.3.4")));
    verify(client).deregisterAll(eq(sourceLoadBalancerId), argThat(set -> set.contains("1.2.3.4")));
    verifyReconcilerIgnore(targetLoadBalancerId, "1.2.3.4");
    verifyReconcilerIgnore(sourceLoadBalancerId, "1.2.3.4");

    // load balancers associated with both source and target jobs are not changed
    verify(client, never()).registerAll(eq(commonLoadBalancerId), any());
    verify(client, never()).deregisterAll(eq(commonLoadBalancerId), any());
    verifyNoReconcilerIgnore(commonLoadBalancerId);
    verifyNoReconcilerIgnore(commonLoadBalancerId);
}
 

@Test
public void movedTaskOnlyTargetAssociatedWithLoadBalancer() {
    String taskId = UUID.randomUUID().toString();
    String sourceJobId = UUID.randomUUID().toString();
    String targetJobId = UUID.randomUUID().toString();
    String targetLoadBalancerId = "lb-" + UUID.randomUUID().toString();
    PublishSubject<JobManagerEvent<?>> taskEvents = PublishSubject.create();

    when(client.registerAll(any(), any())).thenReturn(Completable.complete());
    when(client.deregisterAll(any(), any())).thenReturn(Completable.complete());
    when(v3JobOperations.observeJobs()).thenReturn(taskEvents);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, sourceJobId);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, targetJobId);

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

    assertThat(service.getJobLoadBalancers(sourceJobId).toBlocking().toIterable()).isEmpty();

    assertTrue(service.addLoadBalancer(targetJobId, targetLoadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(targetJobId).toBlocking().toIterable())
            .containsExactlyInAnyOrder(targetLoadBalancerId);

    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();

    Task moved = ServiceJobTask.newBuilder()
            .withJobId(targetJobId)
            .withId(taskId)
            .withStatus(TaskStatus.newBuilder().withState(TaskState.Started).build())
            .withTaskContext(CollectionsExt.asMap(
                    TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "1.2.3.4",
                    TaskAttributes.TASK_ATTRIBUTES_MOVED_FROM_JOB, sourceJobId
            )).build();

    // detect the task is moved and gets registered on the target
    taskEvents.onNext(TaskUpdateEvent.newTaskFromAnotherJob(null, moved, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(1);
    verify(client).registerAll(eq(targetLoadBalancerId), argThat(set -> set.contains("1.2.3.4")));
    verify(client, never()).deregisterAll(any(), any());
    verifyReconcilerIgnore(targetLoadBalancerId, "1.2.3.4");
}
 

@Test
public void movedTaskOnlySourceAssociatedWithLoadBalancer() {
    String taskId = UUID.randomUUID().toString();
    String sourceJobId = UUID.randomUUID().toString();
    String targetJobId = UUID.randomUUID().toString();
    String sourceLoadBalancerId = "lb-" + UUID.randomUUID().toString();
    PublishSubject<JobManagerEvent<?>> taskEvents = PublishSubject.create();

    when(client.registerAll(any(), any())).thenReturn(Completable.complete());
    when(client.deregisterAll(any(), any())).thenReturn(Completable.complete());
    when(v3JobOperations.observeJobs()).thenReturn(taskEvents);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, sourceJobId);
    LoadBalancerTests.applyValidGetJobMock(v3JobOperations, targetJobId);

    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(sourceJobId, sourceLoadBalancerId).await(100, TimeUnit.MILLISECONDS));
    assertThat(service.getJobLoadBalancers(sourceJobId).toBlocking().toIterable())
            .containsExactlyInAnyOrder(sourceLoadBalancerId);

    assertThat(service.getJobLoadBalancers(targetJobId).toBlocking().toIterable()).isEmpty();

    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(0);
    verify(client, never()).registerAll(any(), any());
    verify(client, never()).deregisterAll(any(), any());
    verifyNoReconcilerIgnore();

    Task moved = ServiceJobTask.newBuilder()
            .withJobId(targetJobId)
            .withId(taskId)
            .withStatus(TaskStatus.newBuilder().withState(TaskState.Started).build())
            .withTaskContext(CollectionsExt.asMap(
                    TaskAttributes.TASK_ATTRIBUTES_CONTAINER_IP, "1.2.3.4",
                    TaskAttributes.TASK_ATTRIBUTES_MOVED_FROM_JOB, sourceJobId
            )).build();

    // detect the task is moved and gets deregistered from the source
    taskEvents.onNext(TaskUpdateEvent.newTaskFromAnotherJob(null, moved, callMetadata));
    testScheduler.advanceTimeBy(FLUSH_WAIT_TIME_MS, TimeUnit.MILLISECONDS);

    testSubscriber.assertNoErrors().assertValueCount(1);
    verify(client).deregisterAll(eq(sourceLoadBalancerId), argThat(set -> set.contains("1.2.3.4")));
    verify(client, never()).registerAll(any(), any());
    verifyReconcilerIgnore(sourceLoadBalancerId, "1.2.3.4");
}