/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.streams.processor.internals;

import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.errors.ProcessorStateException;
import org.apache.kafka.streams.errors.StreamsException;
import org.apache.kafka.streams.errors.TaskCorruptedException;
import org.apache.kafka.streams.processor.StandbyUpdateListener;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.internals.StateUpdater.ExceptionAndTask;
import org.apache.kafka.streams.processor.internals.Task.State;
import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;

import org.hamcrest.Matcher;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.mockito.InOrder;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.apache.kafka.common.utils.Utils.mkObjectProperties;
import static org.apache.kafka.streams.StreamsConfig.producerPrefix;
import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.standbyTask;
import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statefulTask;
import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statelessTask;
import static org.apache.kafka.test.StreamsTestUtils.TopologyMetadataBuilder.unnamedTopology;
import static org.apache.kafka.test.TestUtils.waitForCondition;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.not;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyMap;
import static org.mockito.Mockito.atLeast;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

class DefaultStateUpdaterTest {

    private static final int COMMIT_INTERVAL = 100;
    private static final long CALL_TIMEOUT = 1000;
    private static final long VERIFICATION_TIMEOUT = 30000;
    private static final TopicPartition TOPIC_PARTITION_A_0 = new TopicPartition("topicA", 0);
    private static final TopicPartition TOPIC_PARTITION_A_1 = new TopicPartition("topicA", 1);
    private static final TopicPartition TOPIC_PARTITION_B_0 = new TopicPartition("topicB", 0);
    private static final TopicPartition TOPIC_PARTITION_B_1 = new TopicPartition("topicB", 1);
    private static final TopicPartition TOPIC_PARTITION_C_0 = new TopicPartition("topicC", 0);
    private static final TopicPartition TOPIC_PARTITION_D_0 = new TopicPartition("topicD", 0);
    private static final TaskId TASK_0_0 = new TaskId(0, 0);
    private static final TaskId TASK_0_1 = new TaskId(0, 1);
    private static final TaskId TASK_0_2 = new TaskId(0, 2);
    private static final TaskId TASK_1_0 = new TaskId(1, 0);
    private static final TaskId TASK_1_1 = new TaskId(1, 1);
    private static final TaskId TASK_A_0_0 = new TaskId(0, 0, "A");
    private static final TaskId TASK_A_0_1 = new TaskId(0, 1, "A");
    private static final TaskId TASK_B_0_0 = new TaskId(0, 0, "B");
    private static final TaskId TASK_B_0_1 = new TaskId(0, 1, "B");

    // need an auto-tick timer to work for draining with timeout
    private final Time time = new MockTime(1L);
    private final StreamsMetricsImpl metrics = new StreamsMetricsImpl(new Metrics(time), "", time);
    private final StreamsConfig config = new StreamsConfig(configProps(COMMIT_INTERVAL));
    private final ChangelogReader changelogReader = mock(ChangelogReader.class);
    private final TopologyMetadata topologyMetadata = unnamedTopology().build();
    private final DefaultStateUpdater stateUpdater =
        new DefaultStateUpdater("test-state-updater", metrics, config, null, changelogReader, topologyMetadata, time);

    @AfterEach
    public void tearDown() {
        stateUpdater.shutdown(Duration.ofMinutes(1));
    }

    private Properties configProps(final int commitInterval) {
        return mkObjectProperties(mkMap(
            mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, "appId"),
            mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:2171"),
            mkEntry(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2),
            mkEntry(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, commitInterval),
            mkEntry(producerPrefix(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG), commitInterval)
        ));
    }

    @Test
    public void shouldShutdownStateUpdater() throws Exception {
        final StreamTask statelessTask = statelessTask(TASK_0_0).inState(State.RESTORING).build();
        final StreamTask restoredStatefulTask = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StreamTask failedStatefulTask = statefulTask(TASK_1_1, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Set.of(TOPIC_PARTITION_B_0));
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(Set.of(TASK_1_1));
        doThrow(taskCorruptedException).when(changelogReader).restore(mkMap(
            mkEntry(TASK_1_1, failedStatefulTask),
            mkEntry(TASK_0_2, standbyTask)
        ));
        stateUpdater.add(statelessTask);
        stateUpdater.add(restoredStatefulTask);
        stateUpdater.add(failedStatefulTask);
        stateUpdater.add(standbyTask);
        stateUpdater.start();
        verifyRestoredActiveTasks(statelessTask, restoredStatefulTask);
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(taskCorruptedException, failedStatefulTask));
        verifyUpdatingTasks(standbyTask);
        verifyPausedTasks();

        stateUpdater.shutdown(Duration.ofMinutes(1));

        verifyRestoredActiveTasks(statelessTask, restoredStatefulTask);
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(taskCorruptedException, failedStatefulTask));
        verifyUpdatingTasks();
        verifyPausedTasks();
        verify(changelogReader).clear();
    }

    @Test
    public void shouldShutdownStateUpdaterWithPausedTasks() throws Exception {
        final StreamTask statefulTask = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        when(topologyMetadata.isPaused(null)).thenReturn(true);
        stateUpdater.add(statefulTask);
        stateUpdater.add(standbyTask);
        stateUpdater.start();
        verifyRestoredActiveTasks();
        verifyExceptionsAndFailedTasks();
        verifyUpdatingTasks();
        verifyPausedTasks(statefulTask, standbyTask);

        stateUpdater.shutdown(Duration.ofMinutes(1));

        verifyRestoredActiveTasks();
        verifyExceptionsAndFailedTasks();
        verifyUpdatingTasks();
        verifyPausedTasks();
        verify(changelogReader).clear();
    }

    @Test
    public void shouldShutdownStateUpdaterAndRestart() {
        stateUpdater.start();

        stateUpdater.shutdown(Duration.ofMinutes(1));

        stateUpdater.start();

        stateUpdater.shutdown(Duration.ofMinutes(1));

        verify(changelogReader, times(2)).clear();
    }

    @Test
    public void shouldThrowIfRestartedWithNonEmptyRestoredTasks() throws Exception {
        final StreamTask restoredTask = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_A_1));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(restoredTask);
        verifyRestoredActiveTasks(restoredTask);
        stateUpdater.shutdown(Duration.ofMinutes(1));

        final IllegalStateException exception = assertThrows(IllegalStateException.class, stateUpdater::start);

        assertEquals("State updater started with non-empty output queues."
            + " This indicates a bug. Please report at https://issues.apache.org/jira/projects/KAFKA/issues or to the"
            + " dev-mailing list (https://kafka.apache.org/contact).", exception.getMessage());
    }

    @Test
    public void shouldThrowIfRestartedWithNonEmptyFailedTasks() throws Exception {
        final StreamTask failedTask = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(Set.of(TASK_0_0));
        doThrow(taskCorruptedException).when(changelogReader).restore(mkMap(mkEntry(TASK_0_0, failedTask)));
        stateUpdater.start();
        stateUpdater.add(failedTask);
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(taskCorruptedException, failedTask));
        stateUpdater.shutdown(Duration.ofMinutes(1));

        final IllegalStateException exception = assertThrows(IllegalStateException.class, stateUpdater::start);

        assertEquals("State updater started with non-empty output queues."
            + " This indicates a bug. Please report at https://issues.apache.org/jira/projects/KAFKA/issues or to the"
            + " dev-mailing list (https://kafka.apache.org/contact).", exception.getMessage());
    }

    @Test
    public void shouldThrowIfStatelessTaskNotInStateRestoring() {
        shouldThrowIfActiveTaskNotInStateRestoring(statelessTask(TASK_0_0).build());
    }

    @Test
    public void shouldThrowIfStatefulTaskNotInStateRestoring() {
        shouldThrowIfActiveTaskNotInStateRestoring(statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).build());
    }

    private void shouldThrowIfActiveTaskNotInStateRestoring(final StreamTask task) {
        shouldThrowIfTaskNotInGivenState(task, State.RESTORING);
    }

    @Test
    public void shouldThrowIfStandbyTaskNotInStateRunning() {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_B_0)).build();
        shouldThrowIfTaskNotInGivenState(task, State.RUNNING);
    }

    private void shouldThrowIfTaskNotInGivenState(final Task task, final State correctState) {
        for (final State state : State.values()) {
            if (state != correctState) {
                when(task.state()).thenReturn(state);
                assertThrows(IllegalStateException.class, () -> stateUpdater.add(task));
            }
        }
    }

    @Test
    public void shouldThrowIfAddingActiveTasksWithSameId() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask task2 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldThrowIfAddingTasksWithSameId(task1, task2);
    }

    @Test
    public void shouldThrowIfAddingStandbyTasksWithSameId() throws Exception {
        final StandbyTask task1 = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final StandbyTask task2 = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        shouldThrowIfAddingTasksWithSameId(task1, task2);
    }

    @Test
    public void shouldThrowIfAddingActiveAndStandbyTaskWithSameId() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask task2 = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        shouldThrowIfAddingTasksWithSameId(task1, task2);
    }

    @Test
    public void shouldThrowIfAddingStandbyAndActiveTaskWithSameId() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask task2 = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        shouldThrowIfAddingTasksWithSameId(task2, task1);
    }

    private void shouldThrowIfAddingTasksWithSameId(final Task task1, final Task task2) throws Exception {
        stateUpdater.start();
        stateUpdater.add(task1);
        stateUpdater.add(task2);

        verifyFailedTasks(IllegalStateException.class, task1);
        waitForCondition(
            () -> !stateUpdater.isRunning(),
            VERIFICATION_TIMEOUT,
            "Did not switch to non-running within the given timeout!"
        );
    }

    @Test
    public void shouldImmediatelyAddSingleStatelessTaskToRestoredTasks() throws Exception {
        final StreamTask task1 = statelessTask(TASK_0_0).inState(State.RESTORING).build();
        shouldImmediatelyAddStatelessTasksToRestoredTasks(task1);
    }

    @Test
    public void shouldImmediatelyAddMultipleStatelessTasksToRestoredTasks() throws Exception {
        final StreamTask task1 = statelessTask(TASK_0_0).inState(State.RESTORING).build();
        final StreamTask task2 = statelessTask(TASK_0_2).inState(State.RESTORING).build();
        final StreamTask task3 = statelessTask(TASK_1_0).inState(State.RESTORING).build();
        shouldImmediatelyAddStatelessTasksToRestoredTasks(task1, task2, task3);
    }

    private void shouldImmediatelyAddStatelessTasksToRestoredTasks(final StreamTask... tasks) throws Exception {
        stateUpdater.start();
        for (final StreamTask task : tasks) {
            stateUpdater.add(task);
        }

        verifyRestoredActiveTasks(tasks);
        verifyNeverCheckpointTasks(tasks);
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();
    }

    @Test
    public void shouldRestoreSingleActiveStatefulTask() throws Exception {
        final StreamTask task =
            statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final AtomicBoolean allChangelogCompleted = new AtomicBoolean(false);
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet())
            .thenReturn(Set.of(TOPIC_PARTITION_A_0))
            .thenAnswer(invocation -> {
                allChangelogCompleted.set(true);
                return Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0);
            });
        when(changelogReader.allChangelogsCompleted()).thenAnswer(invocation -> allChangelogCompleted.get());
        stateUpdater.start();

        stateUpdater.add(task);

        verifyRestoredActiveTasks(task);
        verifyCheckpointTasks(task);
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();
        verify(changelogReader).register(task.changelogPartitions(), task.stateManager());
        verify(changelogReader).unregister(task.changelogPartitions());
        verify(changelogReader).enforceRestoreActive();
        verify(changelogReader, atLeast(3)).restore(anyMap());
        verify(changelogReader, never()).transitToUpdateStandby();
    }

    @Test
    public void shouldRestoreMultipleActiveStatefulTasks() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask task2 = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StreamTask task3 = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RESTORING).build();
        final AtomicBoolean allChangelogCompleted = new AtomicBoolean(false);
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet())
            .thenReturn(Set.of(TOPIC_PARTITION_C_0))
            .thenReturn(Set.of(TOPIC_PARTITION_C_0, TOPIC_PARTITION_A_0))
            .thenAnswer(invocation -> {
                allChangelogCompleted.set(true);
                return Set.of(TOPIC_PARTITION_C_0, TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0);
            });
        when(changelogReader.allChangelogsCompleted()).thenAnswer(invocation -> allChangelogCompleted.get());
        stateUpdater.start();

        stateUpdater.add(task1);
        stateUpdater.add(task2);
        stateUpdater.add(task3);

        verifyRestoredActiveTasks(task3, task1, task2);
        verifyCheckpointTasks(task3, task1, task2);
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();
        verify(changelogReader).register(task1.changelogPartitions(), task1.stateManager());
        verify(changelogReader).register(task2.changelogPartitions(), task2.stateManager());
        verify(changelogReader).register(task3.changelogPartitions(), task3.stateManager());
        verify(changelogReader).unregister(task1.changelogPartitions());
        verify(changelogReader).unregister(task2.changelogPartitions());
        verify(changelogReader).unregister(task3.changelogPartitions());
        verify(changelogReader, times(3)).enforceRestoreActive();
        verify(changelogReader, atLeast(4)).restore(anyMap());
        verify(changelogReader, never()).transitToUpdateStandby();
    }

    @Test
    public void shouldReturnTrueForRestoreActiveTasksIfTaskAdded() {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RESTORING).build();
        stateUpdater.add(task);

        assertTrue(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldReturnTrueForRestoreActiveTasksIfTaskUpdating() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted())
            .thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyRestoredActiveTasks();
        verifyUpdatingTasks(task);
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();

        assertTrue(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldReturnTrueForRestoreActiveTasksIfTaskRestored() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0));
        when(changelogReader.allChangelogsCompleted())
            .thenReturn(true);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyRestoredActiveTasks(task);
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();

        assertTrue(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldReturnFalseForRestoreActiveTasksIfTaskRemoved() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted())
            .thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED).get();
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();

        assertFalse(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldReturnTrueForRestoreActiveTasksIfTaskFailed() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted())
            .thenReturn(false);
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(Set.of(task.id()));
        doThrow(taskCorruptedException).when(changelogReader).restore(mkMap(mkEntry(TASK_0_0, task)));
        stateUpdater.start();
        stateUpdater.add(task);
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(taskCorruptedException, task));
        verifyPausedTasks();

        assertTrue(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldReturnTrueForRestoreActiveTasksIfTaskPaused() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted())
            .thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyUpdatingTasks(task);
        when(topologyMetadata.isPaused(null)).thenReturn(true);
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks(task);

        assertTrue(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldReturnFalseForRestoreActiveTasksIfTaskRemovedFromStateUpdater() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0));
        when(changelogReader.allChangelogsCompleted())
            .thenReturn(true);
        stateUpdater.start();
        stateUpdater.add(task);
        stateUpdater.drainRestoredActiveTasks(Duration.ofMillis(VERIFICATION_TIMEOUT));
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();

        assertFalse(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldReturnTrueForRestoreActiveTasksIfStandbyTask() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted())
            .thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyRestoredActiveTasks();
        verifyUpdatingTasks(task);
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();

        assertFalse(stateUpdater.restoresActiveTasks());
    }

    @Test
    public void shouldDrainRestoredActiveTasks() throws Exception {
        assertTrue(stateUpdater.drainRestoredActiveTasks(Duration.ZERO).isEmpty());

        final StreamTask task1 = statelessTask(TASK_0_0).inState(State.RESTORING).build();
        stateUpdater.start();
        stateUpdater.add(task1);

        verifyDrainingRestoredActiveTasks(task1);

        final StreamTask task2 = statelessTask(TASK_1_1).inState(State.RESTORING).build();
        final StreamTask task3 = statelessTask(TASK_1_0).inState(State.RESTORING).build();
        final StreamTask task4 = statelessTask(TASK_0_2).inState(State.RESTORING).build();
        stateUpdater.add(task2);
        stateUpdater.add(task3);
        stateUpdater.add(task4);

        verifyDrainingRestoredActiveTasks(task2, task3, task4);
    }

    @Test
    public void shouldUpdateSingleStandbyTask() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0))
            .inState(State.RUNNING).build();
        shouldUpdateStandbyTasks(task);
    }

    @Test
    public void shouldUpdateMultipleStandbyTasks() throws Exception {
        final StandbyTask task1 = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        final StandbyTask task2 = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final StandbyTask task3 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        shouldUpdateStandbyTasks(task1, task2, task3);
    }

    private void shouldUpdateStandbyTasks(final StandbyTask... tasks) throws Exception {
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();

        for (final StandbyTask task : tasks) {
            stateUpdater.add(task);
        }

        verifyUpdatingStandbyTasks(tasks);
        verifyRestoredActiveTasks();
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();
        for (final StandbyTask task : tasks) {
            verify(changelogReader).register(task.changelogPartitions(), task.stateManager());
        }
        verify(changelogReader).transitToUpdateStandby();
        verify(changelogReader, timeout(VERIFICATION_TIMEOUT).atLeast(1)).restore(anyMap());
        verify(changelogReader, never()).enforceRestoreActive();
    }

    @Test
    public void shouldRestoreActiveStatefulTasksAndUpdateStandbyTasks() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask task2 = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask task3 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final StandbyTask task4 = standbyTask(TASK_1_1, Set.of(TOPIC_PARTITION_D_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet())
            .thenReturn(Set.of(TOPIC_PARTITION_A_0))
            .thenReturn(Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();

        stateUpdater.add(task1);
        stateUpdater.add(task2);
        stateUpdater.add(task3);
        stateUpdater.add(task4);

        verifyRestoredActiveTasks(task2, task1);
        verifyCheckpointTasks(task2, task1);
        verifyUpdatingStandbyTasks(task4, task3);
        verifyExceptionsAndFailedTasks();
        verifyPausedTasks();
        verify(changelogReader).register(task1.changelogPartitions(), task1.stateManager());
        verify(changelogReader).register(task2.changelogPartitions(), task2.stateManager());
        verify(changelogReader).register(task3.changelogPartitions(), task3.stateManager());
        verify(changelogReader).register(task4.changelogPartitions(), task4.stateManager());
        verify(changelogReader, atLeast(3)).restore(anyMap());
        final InOrder orderVerifier = inOrder(changelogReader, task1, task2);
        orderVerifier.verify(changelogReader, times(2)).enforceRestoreActive();
        orderVerifier.verify(changelogReader).transitToUpdateStandby();
    }

    @Test
    public void shouldRestoreActiveStatefulTaskThenUpdateStandbyTaskAndAgainRestoreActiveStatefulTask() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask task2 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final StreamTask task3 = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs())
            .thenReturn(Collections.emptySet())
            .thenReturn(Set.of(TOPIC_PARTITION_A_0))
            .thenReturn(Set.of(TOPIC_PARTITION_B_0));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();

        stateUpdater.add(task1);
        stateUpdater.add(task2);

        verifyRestoredActiveTasks(task1);
        verifyCheckpointTasks(task1);
        verifyUpdatingStandbyTasks(task2);
        final InOrder orderVerifier = inOrder(changelogReader);
        orderVerifier.verify(changelogReader, times(1)).enforceRestoreActive();
        orderVerifier.verify(changelogReader, times(1)).transitToUpdateStandby();

        stateUpdater.add(task3);

        verifyRestoredActiveTasks(task1, task3);
        verifyCheckpointTasks(task3);
        orderVerifier.verify(changelogReader, times(1)).enforceRestoreActive();
        orderVerifier.verify(changelogReader, times(1)).transitToUpdateStandby();
    }

    @Test
    public void shouldUpdateStandbyTaskAfterAllActiveStatefulTasksFailed() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final TaskCorruptedException taskCorruptedException =
            new TaskCorruptedException(Set.of(activeTask1.id(), activeTask2.id()));
        final Map<TaskId, Task> updatingTasks1 = mkMap(
            mkEntry(activeTask1.id(), activeTask1),
            mkEntry(activeTask2.id(), activeTask2),
            mkEntry(standbyTask.id(), standbyTask)
        );
        doThrow(taskCorruptedException).doReturn(0L).when(changelogReader).restore(updatingTasks1);
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();

        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask);

        final ExceptionAndTask expectedExceptionAndTask1 = new ExceptionAndTask(taskCorruptedException, activeTask1);
        final ExceptionAndTask expectedExceptionAndTask2 = new ExceptionAndTask(taskCorruptedException, activeTask2);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTask1, expectedExceptionAndTask2);
        final InOrder orderVerifier = inOrder(changelogReader);
        orderVerifier.verify(changelogReader, atLeast(1)).enforceRestoreActive();
        orderVerifier.verify(changelogReader).transitToUpdateStandby();
    }

    @Test
    public void shouldNotTransitToStandbyAgainAfterStandbyTaskFailed() throws Exception {
        final StandbyTask task1 = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        final StandbyTask task2 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final Map<TaskId, Task> updatingTasks = mkMap(
            mkEntry(task1.id(), task1),
            mkEntry(task2.id(), task2)
        );
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(Set.of(task1.id()));
        final ExceptionAndTask expectedExceptionAndTasks = new ExceptionAndTask(taskCorruptedException, task1);
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        doThrow(taskCorruptedException).doReturn(0L).when(changelogReader).restore(updatingTasks);

        stateUpdater.start();
        stateUpdater.add(task1);
        stateUpdater.add(task2);

        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        verify(changelogReader, times(1)).transitToUpdateStandby();
    }

    @Test
    public void shouldUpdateStandbyTaskAfterAllActiveStatefulTasksRemoved() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask);
        verifyUpdatingTasks(activeTask1, activeTask2, standbyTask);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future1 = stateUpdater.remove(activeTask1.id(), StandbyUpdateListener.SuspendReason.MIGRATED);
        final CompletableFuture<StateUpdater.RemovedTaskResult> future2 = stateUpdater.remove(activeTask2.id(), StandbyUpdateListener.SuspendReason.MIGRATED);
        CompletableFuture.allOf(future1, future2).get();

        final InOrder orderVerifier = inOrder(changelogReader);
        orderVerifier.verify(changelogReader, atLeast(1)).enforceRestoreActive();
        orderVerifier.verify(changelogReader).transitToUpdateStandby();
    }

    @Test
    public void shouldNotSwitchTwiceToUpdatingStandbyTaskIfStandbyTaskIsRemoved() throws Exception {
        final StandbyTask standbyTask1 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final StandbyTask standbyTask2 = standbyTask(TASK_0_1, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(standbyTask1);
        stateUpdater.add(standbyTask2);
        verifyUpdatingTasks(standbyTask1, standbyTask2);

        stateUpdater.remove(standbyTask2.id(), StandbyUpdateListener.SuspendReason.MIGRATED).get();

        verify(changelogReader).transitToUpdateStandby();
    }

    @Test
    public void shouldRemoveUpdatingActiveStatefulTask() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldRemoveUpdatingStatefulTask(task);
    }

    @Test
    public void shouldRemoveUpdatingStandbyTask() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        shouldRemoveUpdatingStatefulTask(task);
    }

    @Test
    public void shouldPassSuspendReasonToChangelogReaderOnRemove() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyUpdatingTasks(task);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future =
            stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.PROMOTED);

        assertEquals(new StateUpdater.RemovedTaskResult(task), future.get());
        verify(changelogReader).unregister(task.changelogPartitions(), StandbyUpdateListener.SuspendReason.PROMOTED);
    }

    private void shouldRemoveUpdatingStatefulTask(final Task task) throws Exception {
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyUpdatingTasks(task);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        assertEquals(new StateUpdater.RemovedTaskResult(task), future.get());
        verifyCheckpointTasks(task);
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyPausedTasks();
        verifyExceptionsAndFailedTasks();
        verify(changelogReader).unregister(task.changelogPartitions(), StandbyUpdateListener.SuspendReason.MIGRATED);
    }

    @Test
    public void shouldThrowIfRemovingUpdatingActiveTaskFailsWithStreamsException() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamsException streamsException = new StreamsException("Something happened", task.id());
        setupShouldThrowIfRemovingUpdatingStatefulTaskFailsWithException(task, streamsException);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        verifyRemovingUpdatingStatefulTaskFails(future, task, streamsException, true);

    }

    @Test
    public void shouldThrowIfRemovingUpdatingActiveTaskFailsWithRuntimeException() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final RuntimeException runtimeException = new RuntimeException("Something happened");
        setupShouldThrowIfRemovingUpdatingStatefulTaskFailsWithException(task, runtimeException);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        verifyRemovingUpdatingStatefulTaskFails(future, task, runtimeException, false);
    }

    @Test
    public void shouldThrowIfRemovingUpdatingStandbyTaskFailsWithStreamsException() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        final StreamsException streamsException = new StreamsException("Something happened", task.id());
        setupShouldThrowIfRemovingUpdatingStatefulTaskFailsWithException(task, streamsException);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        verifyRemovingUpdatingStatefulTaskFails(future, task, streamsException, true);
    }

    @Test
    public void shouldThrowIfRemovingUpdatingStandbyTaskFailsWithRuntimeException() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        final RuntimeException runtimeException = new RuntimeException("Something happened");
        setupShouldThrowIfRemovingUpdatingStatefulTaskFailsWithException(task, runtimeException);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        verifyRemovingUpdatingStatefulTaskFails(future, task, runtimeException, false);
    }

    private void setupShouldThrowIfRemovingUpdatingStatefulTaskFailsWithException(final Task task,
                                                                                  final RuntimeException exception) throws Exception {
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        final Collection<TopicPartition> changelogPartitions = task.changelogPartitions();
        doThrow(exception).when(changelogReader).unregister(changelogPartitions, StandbyUpdateListener.SuspendReason.MIGRATED);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyUpdatingTasks(task);
    }

    private void verifyRemovingUpdatingStatefulTaskFails(final CompletableFuture<StateUpdater.RemovedTaskResult> future,
                                                         final Task task,
                                                         final RuntimeException exception,
                                                         final boolean shouldStillBeRunning) throws Exception {
        final ExecutionException executionException = assertThrows(ExecutionException.class, future::get);
        assertInstanceOf(RuntimeException.class, executionException.getCause());
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyPausedTasks();
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(exception, task));
        assertEquals(shouldStillBeRunning, stateUpdater.isRunning());
    }

    @Test
    public void shouldRemovePausedTask() throws Exception {
        final StreamTask statefulTask = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_0_1, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        stateUpdater.start();
        stateUpdater.add(statefulTask);
        stateUpdater.add(standbyTask);
        verifyUpdatingTasks(statefulTask, standbyTask);
        when(topologyMetadata.isPaused(null)).thenReturn(true);
        verifyPausedTasks(statefulTask, standbyTask);
        verifyUpdatingTasks();

        final CompletableFuture<StateUpdater.RemovedTaskResult> futureOfStatefulTask = stateUpdater.remove(statefulTask.id(), StandbyUpdateListener.SuspendReason.MIGRATED);
        final CompletableFuture<StateUpdater.RemovedTaskResult> futureOfStandbyTask = stateUpdater.remove(standbyTask.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        assertEquals(new StateUpdater.RemovedTaskResult(statefulTask), futureOfStatefulTask.get());
        assertEquals(new StateUpdater.RemovedTaskResult(standbyTask), futureOfStandbyTask.get());
        verifyPausedTasks();
        verifyCheckpointTasks(statefulTask, standbyTask);
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
        verify(changelogReader).unregister(statefulTask.changelogPartitions(), StandbyUpdateListener.SuspendReason.MIGRATED);
        verify(changelogReader).unregister(standbyTask.changelogPartitions(), StandbyUpdateListener.SuspendReason.MIGRATED);
    }

    @Test
    public void shouldThrowIfRemovingPausedTaskFails() throws Exception {
        final StreamTask statefulTask = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamsException streamsException = new StreamsException("Something happened", statefulTask.id());
        final Collection<TopicPartition> changelogPartitions = statefulTask.changelogPartitions();
        doThrow(streamsException).when(changelogReader).unregister(changelogPartitions, StandbyUpdateListener.SuspendReason.MIGRATED);
        stateUpdater.start();
        stateUpdater.add(statefulTask);
        verifyUpdatingTasks(statefulTask);
        when(topologyMetadata.isPaused(null)).thenReturn(true);
        verifyPausedTasks(statefulTask);
        verifyUpdatingTasks();

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(statefulTask.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        final ExecutionException executionException = assertThrows(ExecutionException.class, future::get);
        assertInstanceOf(StreamsException.class, executionException.getCause());
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyPausedTasks();
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(streamsException, statefulTask));
        assertTrue(stateUpdater.isRunning());
    }

    @Test
    public void shouldRemoveActiveStatefulTaskFromRestoredActiveTasks() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldRemoveTaskFromRestoredActiveTasks(task);
    }

    @Test
    public void shouldRemoveStatelessTaskFromRestoredActiveTasks() throws Exception {
        final StreamTask task = statelessTask(TASK_0_0).inState(State.RESTORING).build();
        shouldRemoveTaskFromRestoredActiveTasks(task);
    }

    private void shouldRemoveTaskFromRestoredActiveTasks(final StreamTask task) throws Exception {
        when(changelogReader.completedChangelogs()).thenReturn(Collections.singleton(TOPIC_PARTITION_A_0));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        verifyRestoredActiveTasks(task);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED);
        future.get();

        assertEquals(new StateUpdater.RemovedTaskResult(task), future.get());
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyPausedTasks();
        verifyExceptionsAndFailedTasks();
    }

    @Test
    public void shouldRemoveActiveStatefulTaskFromFailedTasks() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldRemoveTaskFromFailedTasks(task);
    }

    @Test
    public void shouldRemoveStandbyTaskFromFailedTasks() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        shouldRemoveTaskFromFailedTasks(task);
    }

    private void shouldRemoveTaskFromFailedTasks(final Task task) throws Exception {
        final StreamsException streamsException = new StreamsException("Something happened", task.id());
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        final Map<TaskId, Task> updatingTasks = mkMap(mkEntry(task.id(), task));
        doThrow(streamsException)
            .doReturn(0L)
            .when(changelogReader).restore(updatingTasks);
        stateUpdater.start();

        stateUpdater.add(task);
        final ExceptionAndTask expectedExceptionAndTasks = new ExceptionAndTask(streamsException, task);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id(), StandbyUpdateListener.SuspendReason.MIGRATED);

        assertEquals(new StateUpdater.RemovedTaskResult(task, streamsException), future.get());
        verifyPausedTasks();
        verifyExceptionsAndFailedTasks();
        verifyUpdatingTasks();
        verifyRestoredActiveTasks();
    }

    @Test
    public void shouldCompleteWithNullIfTaskNotFound() throws Exception {
        final StreamTask updatingTask = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask restoredTask = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StreamTask failedTask = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_C_0)).inState(State.RESTORING).build();
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(Set.of(TASK_0_2));
        doThrow(taskCorruptedException).when(changelogReader).restore(mkMap(
            mkEntry(TASK_0_0, updatingTask),
            mkEntry(TASK_0_2, failedTask)
        ));
        when(changelogReader.completedChangelogs()).thenReturn(Set.of(TOPIC_PARTITION_B_0));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(updatingTask);
        stateUpdater.add(restoredTask);
        stateUpdater.add(failedTask);
        verifyRestoredActiveTasks(restoredTask);
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(taskCorruptedException, failedTask));
        verifyUpdatingTasks(updatingTask);
        verifyPausedTasks();

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(TASK_1_0, StandbyUpdateListener.SuspendReason.MIGRATED);

        assertNull(future.get());
        verifyRestoredActiveTasks(restoredTask);
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(taskCorruptedException, failedTask));
        verifyUpdatingTasks(updatingTask);
        verifyPausedTasks();
    }

    @Test
    public void shouldCompleteWithNullIfNoTasks() throws Exception {
        stateUpdater.start();

        final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(TASK_0_1, StandbyUpdateListener.SuspendReason.MIGRATED);

        assertNull(future.get());
        assertTrue(stateUpdater.isRunning());
    }

    @Test
    public void shouldPauseActiveStatefulTask() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldPauseStatefulTask(task);
        verify(changelogReader, never()).transitToUpdateStandby();
    }

    @Test
    public void shouldPauseStandbyTask() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        shouldPauseStatefulTask(task);
        verify(changelogReader, times(1)).transitToUpdateStandby();
    }

    @Test
    public void shouldPauseActiveTaskAndTransitToUpdateStandby() throws Exception {
        final StreamTask task1 = statefulTask(TASK_A_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask task2 = standbyTask(TASK_B_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();

        when(topologyMetadata.isPaused(task1.id().topologyName())).thenReturn(false).thenReturn(true);

        stateUpdater.start();
        stateUpdater.add(task1);
        stateUpdater.add(task2);

        verifyPausedTasks(task1);
        verifyCheckpointTasks(task1);
        verifyRestoredActiveTasks();
        verifyUpdatingTasks(task2);
        verifyExceptionsAndFailedTasks();
        // shutdown ensures that the test does not end before changelog reader methods verified below are called
        stateUpdater.shutdown(Duration.ofMinutes(1));
        verify(changelogReader, times(1)).enforceRestoreActive();
        verify(changelogReader, times(1)).transitToUpdateStandby();
    }

    @Test
    public void shouldPauseStandbyTaskAndNotTransitToRestoreActive() throws Exception {
        final StandbyTask task1 = standbyTask(TASK_A_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        final StandbyTask task2 = standbyTask(TASK_B_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();

        when(topologyMetadata.isPaused(task1.id().topologyName())).thenReturn(false).thenReturn(true);

        stateUpdater.start();
        stateUpdater.add(task1);
        stateUpdater.add(task2);

        verifyPausedTasks(task1);
        verifyUpdatingTasks(task2);
        verifyCheckpointTasks(task1);
        verify(changelogReader, never()).enforceRestoreActive();
    }

    private void shouldPauseStatefulTask(final Task task) throws Exception {
        stateUpdater.start();
        stateUpdater.add(task);
        verifyUpdatingTasks(task);

        when(topologyMetadata.isPaused(null)).thenReturn(true);

        verifyPausedTasks(task);
        verifyCheckpointTasks(task);
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
    }

    @Test
    public void shouldNotPausingNonExistTasks() throws Exception {
        stateUpdater.start();
        when(topologyMetadata.isPaused(null)).thenReturn(true);

        verifyPausedTasks();
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
    }

    @Test
    public void shouldNotPauseActiveStatefulTaskInRestoredActiveTasks() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask controlTask = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();

        when(changelogReader.completedChangelogs()).thenReturn(Collections.singleton(TOPIC_PARTITION_A_0));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        stateUpdater.add(controlTask);
        verifyRestoredActiveTasks(task);
        verifyUpdatingTasks(controlTask);

        when(topologyMetadata.isPaused(null)).thenReturn(true);

        verifyPausedTasks(controlTask);
        verifyRestoredActiveTasks(task);
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
    }

    @Test
    public void shouldNotPauseActiveStatefulTaskInFailedTasks() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldNotPauseTaskInFailedTasks(task);
    }

    @Test
    public void shouldNotPauseStandbyTaskInFailedTasks() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        shouldNotPauseTaskInFailedTasks(task);
    }

    private void shouldNotPauseTaskInFailedTasks(final Task task) throws Exception {
        final StreamTask controlTask = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();

        final StreamsException streamsException = new StreamsException("Something happened", task.id());
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        final Map<TaskId, Task> updatingTasks = mkMap(
            mkEntry(task.id(), task),
            mkEntry(controlTask.id(), controlTask)
        );
        doThrow(streamsException)
            .doReturn(0L)
            .when(changelogReader).restore(updatingTasks);
        stateUpdater.start();

        stateUpdater.add(task);
        stateUpdater.add(controlTask);
        final ExceptionAndTask expectedExceptionAndTasks = new ExceptionAndTask(streamsException, task);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        verifyUpdatingTasks(controlTask);

        when(topologyMetadata.isPaused(null)).thenReturn(true);

        verifyPausedTasks(controlTask);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        verifyUpdatingTasks();
        verifyRestoredActiveTasks();
    }

    @Test
    public void shouldResumeActiveStatefulTask() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldResumeStatefulTask(task);
        verify(changelogReader, times(2)).enforceRestoreActive();
    }

    @Test
    public void shouldIdleWhenAllTasksPaused() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        stateUpdater.start();
        stateUpdater.add(task);

        when(topologyMetadata.isPaused(null)).thenReturn(true);

        verifyPausedTasks(task);
        verifyIdle();

        when(topologyMetadata.isPaused(null)).thenReturn(false);
        stateUpdater.signalResume();

        verifyPausedTasks();
        verifyUpdatingTasks(task);
    }

    @Test
    public void shouldResumeStandbyTask() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        shouldResumeStatefulTask(task);
        // shutdown ensures that the test does not end before changelog reader methods verified below are called
        stateUpdater.shutdown(Duration.ofMinutes(1));
        verify(changelogReader, times(2)).transitToUpdateStandby();
    }

    private void shouldResumeStatefulTask(final Task task) throws Exception {
        stateUpdater.start();
        stateUpdater.add(task);
        verifyUpdatingTasks(task);

        when(topologyMetadata.isPaused(null)).thenReturn(true);

        verifyPausedTasks(task);
        verifyUpdatingTasks();

        when(topologyMetadata.isPaused(null)).thenReturn(false);
        stateUpdater.signalResume();

        verifyPausedTasks();
        verifyUpdatingTasks(task);
    }

    @Test
    public void shouldNotResumeNonExistingTasks() throws Exception {
        stateUpdater.start();

        verifyPausedTasks();
        verifyRestoredActiveTasks();
        verifyUpdatingTasks();
        verifyExceptionsAndFailedTasks();
    }

    @Test
    public void shouldNotResumeActiveStatefulTaskInRestoredActiveTasks() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask controlTask = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();

        when(changelogReader.completedChangelogs()).thenReturn(Collections.singleton(TOPIC_PARTITION_A_0));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task);
        stateUpdater.add(controlTask);

        verifyRestoredActiveTasks(task);
        verifyPausedTasks();
        verifyRestoredActiveTasks(task);
        verifyUpdatingTasks(controlTask);
        verifyExceptionsAndFailedTasks();
    }

    @Test
    public void shouldNotResumeActiveStatefulTaskInFailedTasks() throws Exception {
        final StreamTask task = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        shouldNotPauseTaskInFailedTasks(task);
    }

    @Test
    public void shouldNotResumeStandbyTaskInFailedTasks() throws Exception {
        final StandbyTask task = standbyTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        shouldNotResumeTaskInFailedTasks(task);
    }

    private void shouldNotResumeTaskInFailedTasks(final Task task) throws Exception {
        final StreamTask controlTask = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StreamsException streamsException = new StreamsException("Something happened", task.id());
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        final Map<TaskId, Task> updatingTasks = mkMap(
            mkEntry(task.id(), task),
            mkEntry(controlTask.id(), controlTask)
        );
        doThrow(streamsException)
            .doReturn(0L)
            .when(changelogReader).restore(updatingTasks);
        stateUpdater.start();

        stateUpdater.add(task);
        stateUpdater.add(controlTask);
        final ExceptionAndTask expectedExceptionAndTasks = new ExceptionAndTask(streamsException, task);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        verifyUpdatingTasks(controlTask);

        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks);
        verifyUpdatingTasks(controlTask);
    }

    @Test
    public void shouldAddFailedTasksToQueueWhenRestoreThrowsStreamsExceptionWithoutTask() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask task2 = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final String exceptionMessage = "The Streams were crossed!";
        final StreamsException streamsException = new StreamsException(exceptionMessage);
        final Map<TaskId, Task> updatingTasks = mkMap(
            mkEntry(task1.id(), task1),
            mkEntry(task2.id(), task2)
        );
        doReturn(0L).doThrow(streamsException).when(changelogReader).restore(updatingTasks);
        stateUpdater.start();

        stateUpdater.add(task1);
        stateUpdater.add(task2);

        final ExceptionAndTask expectedExceptionAndTask1 = new ExceptionAndTask(streamsException, task1);
        final ExceptionAndTask expectedExceptionAndTask2 = new ExceptionAndTask(streamsException, task2);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTask1, expectedExceptionAndTask2);
        verifyPausedTasks();
        verifyUpdatingTasks();
        verifyRestoredActiveTasks();
        assertTrue(stateUpdater.isRunning());
    }

    @Test
    public void shouldAddFailedTasksToQueueWhenRestoreThrowsStreamsExceptionWithTask() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask task2 = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask task3 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final String exceptionMessage = "The Streams were crossed!";
        final StreamsException streamsException1 = new StreamsException(exceptionMessage, task1.id());
        final StreamsException streamsException2 = new StreamsException(exceptionMessage, task3.id());
        final Map<TaskId, Task> updatingTasksBeforeFirstThrow = mkMap(
            mkEntry(task1.id(), task1),
            mkEntry(task2.id(), task2),
            mkEntry(task3.id(), task3)
        );
        final Map<TaskId, Task> updatingTasksBeforeSecondThrow = mkMap(
            mkEntry(task2.id(), task2),
            mkEntry(task3.id(), task3)
        );
        doReturn(0L)
            .doThrow(streamsException1)
            .when(changelogReader).restore(updatingTasksBeforeFirstThrow);
        doReturn(0L)
            .doThrow(streamsException2)
            .when(changelogReader).restore(updatingTasksBeforeSecondThrow);
        stateUpdater.start();

        stateUpdater.add(task1);
        stateUpdater.add(task2);
        stateUpdater.add(task3);

        final ExceptionAndTask expectedExceptionAndTasks1 = new ExceptionAndTask(streamsException1, task1);
        final ExceptionAndTask expectedExceptionAndTasks2 = new ExceptionAndTask(streamsException2, task3);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks1, expectedExceptionAndTasks2);
        verifyUpdatingTasks(task2);
        verifyRestoredActiveTasks();
        verifyPausedTasks();
        assertTrue(stateUpdater.isRunning());
    }

    @Test
    public void shouldHandleTaskCorruptedExceptionAndAddFailedTasksToQueue() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask task2 = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final StreamTask task3 = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RESTORING).build();
        final Set<TaskId> expectedTaskIds = Set.of(task1.id(), task2.id());
        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(expectedTaskIds);
        final Map<TaskId, Task> updatingTasks = mkMap(
            mkEntry(task1.id(), task1),
            mkEntry(task2.id(), task2),
            mkEntry(task3.id(), task3)
        );
        doReturn(0L).doThrow(taskCorruptedException).doReturn(0L).when(changelogReader).restore(updatingTasks);
        stateUpdater.start();

        stateUpdater.add(task1);
        stateUpdater.add(task2);
        stateUpdater.add(task3);

        final ExceptionAndTask expectedExceptionAndTask1 = new ExceptionAndTask(taskCorruptedException, task1);
        final ExceptionAndTask expectedExceptionAndTask2 = new ExceptionAndTask(taskCorruptedException, task2);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTask1, expectedExceptionAndTask2);
        verifyUpdatingTasks(task3);
        verifyRestoredActiveTasks();
        verify(changelogReader).unregister(Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0));
        verify(task1).markChangelogAsCorrupted(Set.of(TOPIC_PARTITION_A_0));
        verify(task2).markChangelogAsCorrupted(Set.of(TOPIC_PARTITION_B_0));
        assertTrue(stateUpdater.isRunning());
    }

    @Test
    public void shouldAddFailedTasksToQueueWhenUncaughtExceptionIsThrown() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask task2 = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RUNNING).build();
        final IllegalStateException illegalStateException = new IllegalStateException("Nobody expects the Spanish inquisition!");
        final Map<TaskId, Task> updatingTasks = mkMap(
            mkEntry(task1.id(), task1),
            mkEntry(task2.id(), task2)
        );
        doThrow(illegalStateException).when(changelogReader).restore(updatingTasks);
        stateUpdater.start();

        stateUpdater.add(task1);
        stateUpdater.add(task2);

        final ExceptionAndTask expectedExceptionAndTask1 = new ExceptionAndTask(illegalStateException, task1);
        final ExceptionAndTask expectedExceptionAndTask2 = new ExceptionAndTask(illegalStateException, task2);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTask1, expectedExceptionAndTask2);
        verifyUpdatingTasks();
        verifyRestoredActiveTasks();
        verifyPausedTasks();
        assertFalse(stateUpdater.isRunning());
    }

    @Test
    public void shouldDrainFailedTasksAndExceptions() throws Exception {
        assertFalse(stateUpdater.hasExceptionsAndFailedTasks());
        assertTrue(stateUpdater.drainExceptionsAndFailedTasks().isEmpty());

        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StreamTask task2 = statefulTask(TASK_1_1, Set.of(TOPIC_PARTITION_C_0)).inState(State.RESTORING).build();
        final StreamTask task3 = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask task4 = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_D_0)).inState(State.RESTORING).build();
        final String exceptionMessage = "The Streams were crossed!";
        final StreamsException streamsException1 = new StreamsException(exceptionMessage, task1.id());
        final Map<TaskId, Task> updatingTasks1 = mkMap(
            mkEntry(task1.id(), task1)
        );
        doThrow(streamsException1)
            .when(changelogReader).restore(updatingTasks1);
        final StreamsException streamsException2 = new StreamsException(exceptionMessage, task2.id());
        final StreamsException streamsException3 = new StreamsException(exceptionMessage, task3.id());
        final StreamsException streamsException4 = new StreamsException(exceptionMessage, task4.id());
        final Map<TaskId, Task> updatingTasks2 = mkMap(
            mkEntry(task2.id(), task2),
            mkEntry(task3.id(), task3),
            mkEntry(task4.id(), task4)
        );
        doThrow(streamsException2).when(changelogReader).restore(updatingTasks2);
        final Map<TaskId, Task> updatingTasks3 = mkMap(
            mkEntry(task3.id(), task3),
            mkEntry(task4.id(), task4)
        );
        doThrow(streamsException3).when(changelogReader).restore(updatingTasks3);
        final Map<TaskId, Task> updatingTasks4 = mkMap(
            mkEntry(task4.id(), task4)
        );
        doThrow(streamsException4).when(changelogReader).restore(updatingTasks4);
        stateUpdater.start();

        stateUpdater.add(task1);

        final ExceptionAndTask expectedExceptionAndTasks1 = new ExceptionAndTask(streamsException1, task1);
        verifyDrainingExceptionsAndFailedTasks(expectedExceptionAndTasks1);

        stateUpdater.add(task2);
        stateUpdater.add(task3);
        stateUpdater.add(task4);

        final ExceptionAndTask expectedExceptionAndTasks2 = new ExceptionAndTask(streamsException2, task2);
        final ExceptionAndTask expectedExceptionAndTasks3 = new ExceptionAndTask(streamsException3, task3);
        final ExceptionAndTask expectedExceptionAndTasks4 = new ExceptionAndTask(streamsException4, task4);
        verifyDrainingExceptionsAndFailedTasks(expectedExceptionAndTasks2, expectedExceptionAndTasks3, expectedExceptionAndTasks4);
    }

    @Test
    public void shouldAutoCheckpointTasksOnInterval() throws Exception {
        final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask task2 = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask task3 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final StandbyTask task4 = standbyTask(TASK_1_1, Set.of(TOPIC_PARTITION_D_0)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(task1);
        stateUpdater.add(task2);
        stateUpdater.add(task3);
        stateUpdater.add(task4);
        // wait for all tasks added to the thread before advance timer
        verifyUpdatingTasks(task1, task2, task3, task4);

        time.sleep(COMMIT_INTERVAL + 1);

        verifyExceptionsAndFailedTasks();
        verifyCheckpointTasks(task1, task2, task3, task4);
    }

    @Test
    public void shouldNotAutoCheckpointTasksIfIntervalNotElapsed() {
        // we need to use a non auto-ticking timer here to control how much time elapsed exactly
        final Time time = new MockTime();
        final DefaultStateUpdater stateUpdater = new DefaultStateUpdater("test-state-updater", metrics, config, null, changelogReader, topologyMetadata, time);
        try {
            final StreamTask task1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
            final StreamTask task2 = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
            final StandbyTask task3 = standbyTask(TASK_1_0, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
            final StandbyTask task4 = standbyTask(TASK_1_1, Set.of(TOPIC_PARTITION_D_0)).inState(State.RUNNING).build();
            when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
            when(changelogReader.allChangelogsCompleted()).thenReturn(false);
            stateUpdater.start();
            stateUpdater.add(task1);
            stateUpdater.add(task2);
            stateUpdater.add(task3);
            stateUpdater.add(task4);

            time.sleep(COMMIT_INTERVAL);

            verifyNeverCheckpointTasks(task1, task2, task3, task4);
        } finally {
            stateUpdater.shutdown(Duration.ofMinutes(1));
        }
    }

    private void verifyCheckpointTasks(final Task... tasks) {
        for (final Task task : tasks) {
            verify(task, timeout(VERIFICATION_TIMEOUT).atLeast(1)).maybeCheckpoint();
        }
    }

    private void verifyNeverCheckpointTasks(final Task... tasks) {
        for (final Task task : tasks) {
            verify(task, never()).maybeCheckpoint();
        }
    }

    @Test
    public void shouldTasksFromInputQueue() {
        stateUpdater.shutdown(Duration.ofMillis(Long.MAX_VALUE));

        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask1 = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final StandbyTask standbyTask2 = standbyTask(TASK_1_1, Set.of(TOPIC_PARTITION_D_0)).inState(State.RUNNING).build();
        final StandbyTask standbyTask3 = standbyTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_1)).inState(State.RUNNING).build();
        stateUpdater.add(activeTask1);
        stateUpdater.add(standbyTask1);
        stateUpdater.add(standbyTask2);
        stateUpdater.remove(TASK_0_0, StandbyUpdateListener.SuspendReason.MIGRATED);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask3);

        verifyGetTasks(Set.of(activeTask1, activeTask2), Set.of(standbyTask1, standbyTask2, standbyTask3));
    }

    @Test
    public void shouldTasks() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask1 = standbyTask(TASK_0_2, Set.of(TOPIC_PARTITION_C_0)).inState(State.RUNNING).build();
        final StandbyTask standbyTask2 = standbyTask(TASK_1_1, Set.of(TOPIC_PARTITION_D_0)).inState(State.RUNNING).build();
        final StandbyTask standbyTask3 = standbyTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_1)).inState(State.RUNNING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(activeTask1);
        stateUpdater.add(standbyTask1);
        stateUpdater.add(standbyTask2);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask3);
        verifyUpdatingTasks(activeTask1, activeTask2, standbyTask1, standbyTask2, standbyTask3);

        verifyGetTasks(Set.of(activeTask1, activeTask2), Set.of(standbyTask1, standbyTask2, standbyTask3));
    }

    @Test
    public void shouldGetTasksFromRestoredActiveTasks() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        when(changelogReader.completedChangelogs()).thenReturn(Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0));
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        stateUpdater.start();
        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        verifyRestoredActiveTasks(activeTask1, activeTask2);

        verifyGetTasks(Set.of(activeTask1, activeTask2), Set.of());

        stateUpdater.drainRestoredActiveTasks(Duration.ofMinutes(1));

        verifyGetTasks(Set.of(), Set.of());
    }

    @Test
    public void shouldGetTasksFromExceptionsAndFailedTasks() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_1_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask2 = standbyTask(TASK_1_1, Set.of(TOPIC_PARTITION_D_0)).inState(State.RUNNING).build();
        final StandbyTask standbyTask1 = standbyTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_1)).inState(State.RUNNING).build();
        final TaskCorruptedException taskCorruptedException =
            new TaskCorruptedException(Set.of(standbyTask1.id(), standbyTask2.id()));
        final StreamsException streamsException = new StreamsException("The Streams were crossed!", activeTask1.id());
        final Map<TaskId, Task> updatingTasks1 = mkMap(
            mkEntry(activeTask1.id(), activeTask1),
            mkEntry(standbyTask1.id(), standbyTask1),
            mkEntry(standbyTask2.id(), standbyTask2)
        );
        doReturn(0L).doThrow(taskCorruptedException).doReturn(0L).when(changelogReader).restore(updatingTasks1);
        final Map<TaskId, Task> updatingTasks2 = mkMap(
            mkEntry(activeTask1.id(), activeTask1)
        );
        doReturn(0L).doThrow(streamsException).doReturn(0L).when(changelogReader).restore(updatingTasks2);
        stateUpdater.start();
        stateUpdater.add(standbyTask1);
        stateUpdater.add(activeTask1);
        stateUpdater.add(standbyTask2);
        final ExceptionAndTask expectedExceptionAndTasks1 = new ExceptionAndTask(taskCorruptedException, standbyTask1);
        final ExceptionAndTask expectedExceptionAndTasks2 = new ExceptionAndTask(taskCorruptedException, standbyTask2);
        final ExceptionAndTask expectedExceptionAndTasks3 = new ExceptionAndTask(streamsException, activeTask1);
        verifyExceptionsAndFailedTasks(expectedExceptionAndTasks1, expectedExceptionAndTasks2, expectedExceptionAndTasks3);

        verifyGetTasks(Set.of(activeTask1), Set.of(standbyTask1, standbyTask2));

        stateUpdater.drainExceptionsAndFailedTasks();

        verifyGetTasks(Set.of(), Set.of());
    }

    @Test
    public void shouldGetTasksFromPausedTasks() throws Exception {
        final StreamTask activeTask = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask = standbyTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_0)).inState(State.RUNNING).build();
        stateUpdater.start();
        stateUpdater.add(activeTask);
        stateUpdater.add(standbyTask);
        verifyUpdatingTasks(activeTask, standbyTask);

        when(topologyMetadata.isPaused(null)).thenReturn(true);

        verifyPausedTasks(activeTask, standbyTask);

        verifyGetTasks(Set.of(activeTask), Set.of(standbyTask));
    }

    @Test
    public void shouldRecordMetrics() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_A_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_B_0_0, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final StandbyTask standbyTask3 = standbyTask(TASK_A_0_1, Set.of(TOPIC_PARTITION_A_1)).inState(State.RUNNING).build();
        final StandbyTask standbyTask4 = standbyTask(TASK_B_0_1, Set.of(TOPIC_PARTITION_B_1)).inState(State.RUNNING).build();
        final Map<TaskId, Task> tasks1234 = mkMap(
            mkEntry(activeTask1.id(), activeTask1),
            mkEntry(activeTask2.id(), activeTask2),
            mkEntry(standbyTask3.id(), standbyTask3),
            mkEntry(standbyTask4.id(), standbyTask4)
        );
        final Map<TaskId, Task> tasks13 = mkMap(
            mkEntry(activeTask1.id(), activeTask1),
            mkEntry(standbyTask3.id(), standbyTask3)
        );

        when(topologyMetadata.isPaused(activeTask2.id().topologyName())).thenReturn(true);
        when(topologyMetadata.isPaused(standbyTask4.id().topologyName())).thenReturn(true);
        when(changelogReader.completedChangelogs()).thenReturn(Collections.emptySet());
        when(changelogReader.allChangelogsCompleted()).thenReturn(false);
        when(changelogReader.restore(tasks1234)).thenReturn(1L);
        when(changelogReader.restore(tasks13)).thenReturn(1L);
        when(changelogReader.isRestoringActive()).thenReturn(true);

        stateUpdater.start();
        stateUpdater.add(activeTask1);
        stateUpdater.add(activeTask2);
        stateUpdater.add(standbyTask3);
        stateUpdater.add(standbyTask4);

        verifyPausedTasks(activeTask2, standbyTask4);
        assertThat(metrics.metrics().size(), is(11));

        final Map<String, String> tagMap = new LinkedHashMap<>();
        tagMap.put("thread-id", "test-state-updater");

        MetricName metricName = new MetricName("active-restoring-tasks",
            "stream-state-updater-metrics",
            "The number of active tasks currently undergoing restoration",
            tagMap);
        verifyMetric(metrics, metricName, is(1.0));

        metricName = new MetricName("standby-updating-tasks",
            "stream-state-updater-metrics",
            "The number of standby tasks currently undergoing state update",
            tagMap);
        verifyMetric(metrics, metricName, is(1.0));

        metricName = new MetricName("active-paused-tasks",
            "stream-state-updater-metrics",
            "The number of active tasks paused restoring",
            tagMap);
        verifyMetric(metrics, metricName, is(1.0));

        metricName = new MetricName("standby-paused-tasks",
            "stream-state-updater-metrics",
            "The number of standby tasks paused state update",
            tagMap);
        verifyMetric(metrics, metricName, is(1.0));

        metricName = new MetricName("idle-ratio",
            "stream-state-updater-metrics",
            "The ratio, over a rolling measurement window, of the time this thread spent being idle",
            tagMap);
        verifyMetric(metrics, metricName, greaterThanOrEqualTo(0.0d));

        metricName = new MetricName("active-restore-ratio",
            "stream-state-updater-metrics",
            "The ratio, over a rolling measurement window, of the time this thread spent restoring active tasks",
            tagMap);
        verifyMetric(metrics, metricName, greaterThanOrEqualTo(0.0d));

        metricName = new MetricName("standby-update-ratio",
            "stream-state-updater-metrics",
            "The ratio, over a rolling measurement window, of the time this thread spent updating standby tasks",
            tagMap);
        verifyMetric(metrics, metricName, is(0.0d));

        metricName = new MetricName("checkpoint-ratio",
            "stream-state-updater-metrics",
            "The ratio, over a rolling measurement window, of the time this thread spent checkpointing tasks restored progress",
            tagMap);
        verifyMetric(metrics, metricName, greaterThanOrEqualTo(0.0d));

        metricName = new MetricName("restore-records-rate",
            "stream-state-updater-metrics",
            "The average per-second number of records restored",
            tagMap);
        verifyMetric(metrics, metricName, not(0.0d));

        metricName = new MetricName("restore-call-rate",
            "stream-state-updater-metrics",
            "The average per-second number of restore calls triggered",
            tagMap);
        verifyMetric(metrics, metricName, not(0.0d));

        stateUpdater.shutdown(Duration.ofMinutes(1));
        assertThat(metrics.metrics().size(), is(1));
    }

    @Test
    public void shouldRemoveMetricsWithoutInterference() {
        final DefaultStateUpdater stateUpdater2 =
            new DefaultStateUpdater("test-state-updater2", metrics, config, null, changelogReader, topologyMetadata, time);
        final List<MetricName> threadMetrics = getMetricNames("test-state-updater");
        final List<MetricName> threadMetrics2 = getMetricNames("test-state-updater2");

        stateUpdater.start();
        stateUpdater2.start();

        for (final MetricName metricName : threadMetrics) {
            assertTrue(metrics.metrics().containsKey(metricName));
        }
        for (final MetricName metricName : threadMetrics2) {
            assertTrue(metrics.metrics().containsKey(metricName));
        }

        stateUpdater2.shutdown(Duration.ofMinutes(1));

        for (final MetricName metricName : threadMetrics) {
            assertTrue(metrics.metrics().containsKey(metricName));
        }
        for (final MetricName metricName : threadMetrics2) {
            assertFalse(metrics.metrics().containsKey(metricName));
        }

        stateUpdater.shutdown(Duration.ofMinutes(1));

        for (final MetricName metricName : threadMetrics) {
            assertFalse(metrics.metrics().containsKey(metricName));
        }
        for (final MetricName metricName : threadMetrics2) {
            assertFalse(metrics.metrics().containsKey(metricName));
        }
    }

    @Test
    public void shouldNotFailTheThreadIfMaybeCheckpointFails() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask failedStatefulTask = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final ProcessorStateException processorStateException = new ProcessorStateException("flush");
        doThrow(processorStateException).when(failedStatefulTask).maybeCheckpoint();

        stateUpdater.add(failedStatefulTask);
        stateUpdater.add(activeTask1);
        stateUpdater.start();
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(processorStateException, failedStatefulTask));
        verifyUpdatingTasks(activeTask1);

        stateUpdater.add(activeTask2);
        verifyUpdatingTasks(activeTask1, activeTask2);
    }

    @Test
    public void shouldNotFailTheThreadIfMaybeCheckpointFailsForCorruptedTask() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask failedStatefulTask = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final ProcessorStateException processorStateException = new ProcessorStateException("flush");
        doThrow(processorStateException).when(failedStatefulTask).maybeCheckpoint();

        final TaskCorruptedException taskCorruptedException = new TaskCorruptedException(Set.of(TASK_0_2));
        when(changelogReader.restore(Map.of(
                TASK_0_0, activeTask1,
                TASK_0_2, failedStatefulTask))
        ).thenThrow(taskCorruptedException);

        stateUpdater.add(failedStatefulTask);
        stateUpdater.add(activeTask1);
        stateUpdater.start();
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(taskCorruptedException, failedStatefulTask));
        verifyUpdatingTasks(activeTask1);

        stateUpdater.add(activeTask2);
        verifyUpdatingTasks(activeTask1, activeTask2);
    }

    @Test
    public void shouldNotFailTheThreadIfMaybeCheckpointFailsDuringTaskRemoval() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask failedStatefulTask = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final ProcessorStateException processorStateException = new ProcessorStateException("flush");
        final AtomicBoolean throwException = new AtomicBoolean(false);
        doAnswer(invocation -> {
            if (throwException.get()) {
                throw processorStateException;
            }
            return null;
        }).when(failedStatefulTask).maybeCheckpoint();
        when(changelogReader.allChangelogsCompleted()).thenReturn(true);

        stateUpdater.add(failedStatefulTask);
        stateUpdater.add(activeTask1);
        stateUpdater.start();
        verifyUpdatingTasks(failedStatefulTask, activeTask1);

        throwException.set(true);
        final ExecutionException exception = assertThrows(ExecutionException.class, () -> stateUpdater.remove(TASK_0_2, StandbyUpdateListener.SuspendReason.MIGRATED).get());
        assertEquals(processorStateException, exception.getCause());

        stateUpdater.add(activeTask2);
        verifyUpdatingTasks(activeTask1, activeTask2);
    }

    @Test
    public void shouldNotFailTheThreadIfMaybeCheckpointFailsDuringTaskPause() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask failedStatefulTask = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final ProcessorStateException processorStateException = new ProcessorStateException("flush");
        doThrow(processorStateException).when(failedStatefulTask).maybeCheckpoint();
        when(topologyMetadata.isPaused(null)).thenReturn(false).thenReturn(false).thenReturn(true);

        stateUpdater.add(failedStatefulTask);
        stateUpdater.add(activeTask1);
        stateUpdater.start();
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(processorStateException, failedStatefulTask));
        verifyPausedTasks(activeTask1);

        stateUpdater.add(activeTask2);
        verifyPausedTasks(activeTask1, activeTask2);
    }

    @Test
    public void shouldNotFailTheThreadIfMaybeCheckpointFailsDuringTaskRestore() throws Exception {
        final StreamTask activeTask1 = statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask activeTask2 = statefulTask(TASK_0_1, Set.of(TOPIC_PARTITION_A_0)).inState(State.RESTORING).build();
        final StreamTask failedStatefulTask = statefulTask(TASK_0_2, Set.of(TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
        final ProcessorStateException processorStateException = new ProcessorStateException("flush");
        doThrow(processorStateException).when(failedStatefulTask).maybeCheckpoint();
        when(changelogReader.completedChangelogs()).thenReturn(Set.of(TOPIC_PARTITION_B_0));

        stateUpdater.add(failedStatefulTask);
        stateUpdater.add(activeTask1);
        stateUpdater.start();
        verifyExceptionsAndFailedTasks(new ExceptionAndTask(processorStateException, failedStatefulTask));
        verifyUpdatingTasks(activeTask1);

        stateUpdater.add(activeTask2);
        verifyUpdatingTasks(activeTask1, activeTask2);
    }

    private static List<MetricName> getMetricNames(final String threadId) {
        final Map<String, String> tagMap = Map.of("thread-id", threadId);
        return List.of(
            new MetricName("active-restoring-tasks", "stream-state-updater-metrics", "", tagMap),
            new MetricName("standby-updating-tasks", "stream-state-updater-metrics", "", tagMap),
            new MetricName("active-paused-tasks", "stream-state-updater-metrics", "", tagMap),
            new MetricName("standby-paused-tasks", "stream-state-updater-metrics", "", tagMap),
            new MetricName("idle-ratio", "stream-state-updater-metrics", "", tagMap),
            new MetricName("standby-update-ratio", "stream-state-updater-metrics", "", tagMap),
            new MetricName("checkpoint-ratio", "stream-state-updater-metrics", "", tagMap),
            new MetricName("restore-records-rate", "stream-state-updater-metrics", "", tagMap),
            new MetricName("restore-call-rate", "stream-state-updater-metrics", "", tagMap)
        );
    }

    @SuppressWarnings("unchecked")
    private static <T> void verifyMetric(final StreamsMetricsImpl metrics,
                                         final MetricName metricName,
                                         final Matcher<T> matcher) {
        assertThat(metrics.metrics().get(metricName).metricName().description(), is(metricName.description()));
        assertThat(metrics.metrics().get(metricName).metricName().tags(), is(metricName.tags()));
        assertThat((T) metrics.metrics().get(metricName).metricValue(), matcher);
    }

    private void verifyGetTasks(final Set<StreamTask> expectedActiveTasks,
                                final Set<StandbyTask> expectedStandbyTasks) {
        final Set<Task> tasks = stateUpdater.tasks();

        assertEquals(expectedActiveTasks.size() + expectedStandbyTasks.size(), tasks.size());
        tasks.forEach(task -> assertInstanceOf(ReadOnlyTask.class, task));
        final Set<TaskId> actualTaskIds = tasks.stream().map(Task::id).collect(Collectors.toSet());
        final Set<Task> expectedTasks = new HashSet<>(expectedActiveTasks);
        expectedTasks.addAll(expectedStandbyTasks);
        final Set<TaskId> expectedTaskIds = expectedTasks.stream().map(Task::id).collect(Collectors.toSet());
        assertTrue(actualTaskIds.containsAll(expectedTaskIds));

        final Set<StreamTask> activeTasks = stateUpdater.activeTasks();
        assertEquals(expectedActiveTasks.size(), activeTasks.size());
        assertTrue(activeTasks.containsAll(expectedActiveTasks));

        final Set<StandbyTask> standbyTasks = stateUpdater.standbyTasks();
        assertEquals(expectedStandbyTasks.size(), standbyTasks.size());
        assertTrue(standbyTasks.containsAll(expectedStandbyTasks));
    }

    private void verifyRestoredActiveTasks(final StreamTask... tasks) throws Exception {
        if (tasks.length == 0) {
            waitForCondition(
                () -> stateUpdater.restoredActiveTasks().isEmpty(),
                VERIFICATION_TIMEOUT,
                "Did not get empty restored active task within the given timeout!"
            );
        } else {
            final Set<StreamTask> expectedRestoredTasks = Set.of(tasks);
            final Set<StreamTask> restoredTasks = new HashSet<>();
            waitForCondition(
                () -> {
                    restoredTasks.addAll(stateUpdater.restoredActiveTasks());
                    return restoredTasks.containsAll(expectedRestoredTasks)
                        && restoredTasks.size() == expectedRestoredTasks.size();
                },
                VERIFICATION_TIMEOUT,
                () -> "Did not get all restored active task within the given timeout! Expected: "
                        + expectedRestoredTasks + ", actual: " + restoredTasks
            );
        }
    }

    private void verifyDrainingRestoredActiveTasks(final StreamTask... tasks) throws Exception {
        final Set<StreamTask> expectedRestoredTasks = Set.of(tasks);
        final Set<StreamTask> restoredTasks = new HashSet<>();
        waitForCondition(
            () -> {
                restoredTasks.addAll(stateUpdater.drainRestoredActiveTasks(Duration.ofMillis(CALL_TIMEOUT)));
                return restoredTasks.containsAll(expectedRestoredTasks)
                    && restoredTasks.size() == expectedRestoredTasks.size();
            },
            VERIFICATION_TIMEOUT,
            () -> "Did not get all restored active task within the given timeout! Expected: "
                    + expectedRestoredTasks + ", actual: " + restoredTasks
        );
        assertTrue(stateUpdater.drainRestoredActiveTasks(Duration.ZERO).isEmpty());
    }

    private void verifyUpdatingTasks(final Task... tasks) throws Exception {
        if (tasks.length == 0) {
            waitForCondition(
                () -> stateUpdater.updatingTasks().isEmpty(),
                VERIFICATION_TIMEOUT,
                "Did not get empty updating task within the given timeout!"
            );
        } else {
            final Set<Task> expectedUpdatingTasks = Set.of(tasks);
            final Set<Task> updatingTasks = new HashSet<>();
            waitForCondition(
                () -> {
                    updatingTasks.addAll(stateUpdater.updatingTasks());
                    return updatingTasks.containsAll(expectedUpdatingTasks)
                        && updatingTasks.size() == expectedUpdatingTasks.size();
                },
                VERIFICATION_TIMEOUT,
                () -> "Did not get all updating task within the given timeout! Expected: "
                        + expectedUpdatingTasks + ", actual: " + updatingTasks
            );
        }
    }

    private void verifyUpdatingStandbyTasks(final StandbyTask... tasks) throws Exception {
        final Set<StandbyTask> expectedStandbyTasks = Set.of(tasks);
        final Set<StandbyTask> standbyTasks = new HashSet<>();
        waitForCondition(
            () -> {
                standbyTasks.addAll(stateUpdater.updatingStandbyTasks());
                return standbyTasks.containsAll(expectedStandbyTasks)
                    && standbyTasks.size() == expectedStandbyTasks.size();
            },
            VERIFICATION_TIMEOUT,
            () -> "Did not see all standby task within the given timeout! Expected: "
                    + expectedStandbyTasks + ", actual: " + standbyTasks
        );
    }

    private void verifyIdle() throws Exception {
        waitForCondition(
            stateUpdater::isIdle,
            VERIFICATION_TIMEOUT,
            "State updater did not enter an idling state!"
        );
    }

    private void verifyPausedTasks(final Task... tasks) throws Exception {
        if (tasks.length == 0) {
            waitForCondition(
                () -> stateUpdater.pausedTasks().isEmpty(),
                VERIFICATION_TIMEOUT,
                "Did not get empty paused task within the given timeout!"
            );
        } else {
            final Set<Task> expectedPausedTasks = Set.of(tasks);
            final Set<Task> pausedTasks = new HashSet<>();
            waitForCondition(
                () -> {
                    pausedTasks.addAll(stateUpdater.pausedTasks());
                    return pausedTasks.containsAll(expectedPausedTasks)
                        && pausedTasks.size() == expectedPausedTasks.size();
                },
                VERIFICATION_TIMEOUT,
                () -> "Did not get all paused task within the given timeout! Expected: "
                        + expectedPausedTasks + ", actual: " + pausedTasks
            );
        }
    }

    private void verifyExceptionsAndFailedTasks(final ExceptionAndTask... exceptionsAndTasks) throws Exception {
        final List<ExceptionAndTask> expectedExceptionAndTasks = Arrays.asList(exceptionsAndTasks);
        final Set<ExceptionAndTask> failedTasks = new HashSet<>();
        waitForCondition(
            () -> {
                failedTasks.addAll(stateUpdater.exceptionsAndFailedTasks());
                return failedTasks.containsAll(expectedExceptionAndTasks)
                    && failedTasks.size() == expectedExceptionAndTasks.size();
            },
            VERIFICATION_TIMEOUT,
            () -> "Did not get all exceptions and failed tasks within the given timeout! Expected: "
                    + expectedExceptionAndTasks + ", actual: " + failedTasks
        );
    }

    private void verifyFailedTasks(final Class<? extends RuntimeException> clazz, final Task... tasks) throws Exception {
        final List<Task> expectedFailedTasks = Arrays.asList(tasks);
        final Set<Task> failedTasks = new HashSet<>();
        waitForCondition(
            () -> {
                for (final ExceptionAndTask exceptionAndTask : stateUpdater.exceptionsAndFailedTasks()) {
                    if (clazz.isInstance(exceptionAndTask.exception())) {
                        failedTasks.add(exceptionAndTask.task());
                    }
                }
                return failedTasks.containsAll(expectedFailedTasks)
                    && failedTasks.size() == expectedFailedTasks.size();
            },
            VERIFICATION_TIMEOUT,
            () -> "Did not get all exceptions and failed tasks within the given timeout! Expected: "
                        + expectedFailedTasks + ", actual: " + failedTasks
        );
    }

    private void verifyDrainingExceptionsAndFailedTasks(final ExceptionAndTask... exceptionsAndTasks) throws Exception {
        final List<ExceptionAndTask> expectedExceptionAndTasks = Arrays.asList(exceptionsAndTasks);
        final List<ExceptionAndTask> failedTasks = new ArrayList<>();
        waitForCondition(
            () -> {
                if (stateUpdater.hasExceptionsAndFailedTasks()) {
                    final List<ExceptionAndTask> exceptionAndTasks = stateUpdater.drainExceptionsAndFailedTasks();
                    assertFalse(exceptionAndTasks.isEmpty());
                    failedTasks.addAll(exceptionAndTasks);
                }
                return failedTasks.containsAll(expectedExceptionAndTasks)
                    && failedTasks.size() == expectedExceptionAndTasks.size();
            },
            VERIFICATION_TIMEOUT,
            () -> "Did not get all exceptions and failed tasks within the given timeout! Expected: "
                    + expectedExceptionAndTasks + ", actual: " + failedTasks
        );
        assertFalse(stateUpdater.hasExceptionsAndFailedTasks());
        assertTrue(stateUpdater.drainExceptionsAndFailedTasks().isEmpty());
    }
}
