Source:
public void shouldReturnAllActiveTasksToPreviousOwnerRegardlessOfBalanceAndTriggerRebalanceIfEndOffsetFetchFailsAndHighAvailabilityEnabled() { builder.addSource(null, "source1", null, null, null, "topic1"); builder.addProcessor("processor1", new MockProcessorSupplier<>(), "source1"); builder.addStateStore(new MockKeyValueStoreBuilder("store1", false), "processor1"); final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2); createMockTaskManager(allTasks); adminClient = EasyMock.createMock(AdminClient.class); expect(adminClient.listOffsets(anyObject())).andThrow(new StreamsException("Should be handled")); configureDefaultPartitionAssignor(); final String firstConsumer = "consumer1"; final String newConsumer = "consumer2"; subscriptions.put(firstConsumer, new Subscription( singletonList("source1"), getInfo(UUID_1, allTasks).encode() )); subscriptions.put(newConsumer, new Subscription( singletonList("source1"), getInfo(UUID_2, EMPTY_TASKS).encode() )); final Map<String, Assignment> assignments = partitionAssignor .assign(metadata, new GroupSubscription(subscriptions)) .groupAssignment(); final List<TaskId> firstConsumerActiveTasks = AssignmentInfo.decode(assignments.get(firstConsumer).userData()).activeTasks(); final List<TaskId> newConsumerActiveTasks = AssignmentInfo.decode(assignments.get(newConsumer).userData()).activeTasks(); assertThat(firstConsumerActiveTasks, equalTo(new ArrayList<>(allTasks))); assertTrue(newConsumerActiveTasks.isEmpty()); assertThat(assignmentError.get(), equalTo(AssignorError.REBALANCE_NEEDED.code())); }

Target:
public void shouldReturnAllActiveTasksToPreviousOwnerRegardlessOfBalanceAndTriggerRebalanceIfEndOffsetFetchFailsAndHighAvailabilityEnabled() { final long rebalanceInterval = 5 * 60 * 1000L; builder.addSource(null, "source1", null, null, null, "topic1"); builder.addProcessor("processor1", new MockProcessorSupplier<>(), "source1"); builder.addStateStore(new MockKeyValueStoreBuilder("store1", false), "processor1"); final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2); createMockTaskManager(allTasks); adminClient = EasyMock.createMock(AdminClient.class); expect(adminClient.listOffsets(anyObject())).andThrow(new StreamsException("Should be handled")); configurePartitionAssignorWith(singletonMap(StreamsConfig.PROBING_REBALANCE_INTERVAL_MS_CONFIG, rebalanceInterval)); final String firstConsumer = "consumer1"; final String newConsumer = "consumer2"; subscriptions.put(firstConsumer, new Subscription( singletonList("source1"), getInfo(UUID_1, allTasks).encode() )); subscriptions.put(newConsumer, new Subscription( singletonList("source1"), getInfo(UUID_2, EMPTY_TASKS).encode() )); final Map<String, Assignment> assignments = partitionAssignor .assign(metadata, new GroupSubscription(subscriptions)) .groupAssignment(); final AssignmentInfo firstConsumerUserData = AssignmentInfo.decode(assignments.get(firstConsumer).userData()); final List<TaskId> firstConsumerActiveTasks = firstConsumerUserData.activeTasks(); final AssignmentInfo newConsumerUserData = AssignmentInfo.decode(assignments.get(newConsumer).userData()); final List<TaskId> newConsumerActiveTasks = newConsumerUserData.activeTasks(); assertThat(firstConsumerActiveTasks, equalTo(new ArrayList<>(allTasks))); assertThat(newConsumerActiveTasks, empty()); assertThat( time.milliseconds() + rebalanceInterval, anyOf( is(firstConsumerUserData.nextRebalanceMs()), is(newConsumerUserData.nextRebalanceMs()) ) ); }