diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java b/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java
index e9e6c056bf63a..e5b47da0c4587 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java
@@ -67,7 +67,7 @@ public class MockProducer<K, V> implements Producer<K, V> {
     private boolean transactionCommitted;
     private boolean transactionAborted;
     private boolean producerFenced;
-    private boolean producerFencedOnClose;
+    private boolean producerFencedOnCommitTxn;
     private boolean sentOffsets;
     private long commitCount = 0L;
     private Map<MetricName, Metric> mockMetrics;
@@ -182,6 +182,10 @@ public void sendOffsetsToTransaction(Map<TopicPartition, OffsetAndMetadata> offs
 
     @Override
     public void commitTransaction() throws ProducerFencedException {
+        if (producerFencedOnCommitTxn) {
+            throw new ProducerFencedException("Producer is fenced");
+        }
+
         verifyProducerState();
         verifyTransactionsInitialized();
         verifyNoTransactionInFlight();
@@ -325,9 +329,6 @@ public void close() {
 
     @Override
     public void close(Duration timeout) {
-        if (producerFencedOnClose) {
-            throw new ProducerFencedException("MockProducer is fenced.");
-        }
         this.closed = true;
     }
 
@@ -341,10 +342,10 @@ public synchronized void fenceProducer() {
         this.producerFenced = true;
     }
 
-    public void fenceProducerOnClose() {
+    public void fenceProducerOnCommitTxn() {
         verifyProducerState();
         verifyTransactionsInitialized();
-        this.producerFencedOnClose = true;
+        this.producerFencedOnCommitTxn = true;
     }
 
     public boolean transactionInitialized() {
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java
index 370106fa74fdb..24258f957f6bf 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilder.java
@@ -1226,14 +1226,15 @@ synchronized Collection<String> sourceTopicCollection() {
     }
 
     synchronized Pattern sourceTopicPattern() {
-        log.debug("Found pattern subscribed source topics, falling back to pattern subscription for the main consumer.");
-
         if (topicPattern == null) {
             final List<String> allSourceTopics = maybeDecorateInternalSourceTopics(sourceTopicNames);
             Collections.sort(allSourceTopics);
             topicPattern = buildPattern(allSourceTopics, nodeToSourcePatterns.values());
         }
 
+        log.debug("Found pattern subscribed source topics, falling back to pattern " +
+            "subscription for the main consumer: {}", topicPattern);
+
         return topicPattern;
     }
 
@@ -1870,19 +1871,7 @@ private boolean hasSubscriptionUpdates() {
         return !subscriptionUpdates.isEmpty();
     }
 
-    synchronized void updateSubscribedTopics(final Set<String> topics,
-                                             final String logPrefix) {
-        log.debug("{}found {} topics possibly matching subscription", logPrefix, topics);
-        subscriptionUpdates.clear();
-        subscriptionUpdates.addAll(topics);
-
-        log.debug("{}updating builder with {} topic(s) with possible matching regex subscription(s)",
-            logPrefix, subscriptionUpdates);
-        setRegexMatchedTopicsToSourceNodes();
-        setRegexMatchedTopicToStateStore();
-    }
-
-    void addSubscribedTopicsFromAssignment(final List<TopicPartition> partitions, final String logPrefix) {
+    synchronized void addSubscribedTopicsFromAssignment(final List<TopicPartition> partitions, final String logPrefix) {
         if (sourceTopicPattern() != null) {
             final Set<String> assignedTopics = new HashSet<>();
             for (final TopicPartition topicPartition : partitions) {
@@ -1892,21 +1881,33 @@ void addSubscribedTopicsFromAssignment(final List<TopicPartition> partitions, fi
             final Collection<String> existingTopics = subscriptionUpdates();
             if (!existingTopics.containsAll(assignedTopics)) {
                 assignedTopics.addAll(existingTopics);
+                updateSubscribedTopics(assignedTopics, logPrefix);
+            }
+        }
+    }
 
-                log.debug("{}found {} topics possibly matching subscription", logPrefix, assignedTopics);
-                subscriptionUpdates.clear();
-                subscriptionUpdates.addAll(assignedTopics);
-
-                log.debug("{}updating builder with {} topic(s) with possible matching regex subscription(s)",
-                    logPrefix, subscriptionUpdates);
-                setRegexMatchedTopicsToSourceNodes();
-                setRegexMatchedTopicToStateStore();
+    synchronized void addSubscribedTopicsFromMetadata(final Set<String> topics, final String logPrefix) {
+        if (sourceTopicPattern() != null) {
+            final Collection<String> existingTopics = subscriptionUpdates();
+            if (!existingTopics.equals(topics)) {
+                topics.addAll(existingTopics);
+                updateSubscribedTopics(topics, logPrefix);
             }
         }
     }
 
-    // following functions are for test only
+    private void updateSubscribedTopics(final Set<String> topics, final String logPrefix) {
+        log.debug("{}found {} topics possibly matching subscription", logPrefix, topics);
+        subscriptionUpdates.clear();
+        subscriptionUpdates.addAll(topics);
 
+        log.debug("{}updating builder with {} topic(s) with possible matching regex subscription(s)",
+            logPrefix, subscriptionUpdates);
+        setRegexMatchedTopicsToSourceNodes();
+        setRegexMatchedTopicToStateStore();
+    }
+
+    // following functions are for test only
     public synchronized Set<String> sourceTopicNames() {
         return sourceTopicNames;
     }
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorStateManager.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorStateManager.java
index 56257d68d107d..98a47b486f9ce 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorStateManager.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorStateManager.java
@@ -189,7 +189,8 @@ public StateStore getGlobalStore(final String name) {
         return globalStores.get(name);
     }
 
-    public void initializeStoreOffsetsFromCheckpoint() {
+    // package-private for test only
+    void initializeStoreOffsetsFromCheckpoint() {
         try {
             final Map<TopicPartition, Long> loadedCheckpoints = checkpointFile.read();
 
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollector.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollector.java
index 7b708ce6f2b8e..78724f62d7ed3 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollector.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollector.java
@@ -27,11 +27,6 @@
 
 public interface RecordCollector extends AutoCloseable {
 
-    /**
-     * Initialize the record collector
-     */
-    void initialize();
-
     <K, V> void send(final String topic,
                      final K key,
                      final V value,
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
index 2a8e532b627b7..a6399858fa74a 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
@@ -73,6 +73,9 @@ public class RecordCollectorImpl implements RecordCollector {
     private Producer<byte[], byte[]> producer;
     private volatile KafkaException sendException;
 
+    /**
+     * @throws StreamsException fatal error that should cause the thread to die (from producer.initTxn)
+     */
     public RecordCollectorImpl(final TaskId taskId,
                                final StreamsConfig config,
                                final LogContext logContext,
@@ -93,16 +96,6 @@ public RecordCollectorImpl(final TaskId taskId,
 
         producer = producerSupplier.get(taskId);
 
-        // TODO K9113: this should be moved to task when it transits to running from created / restoring
-        // then even if there's a long time between the initialization and the first txn that is also fine.
-        initialize();
-    }
-
-    /**
-     * @throws StreamsException fatal error that should cause the thread to die
-     */
-    @Override
-    public void initialize() {
         maybeInitTxns();
     }
 
@@ -333,7 +326,9 @@ private void checkForException() {
     @Override
     public void flush() {
         log.debug("Flushing record collector");
+
         producer.flush();
+
         checkForException();
     }
 
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java
index 11db7417769e2..c69fe42a33bce 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java
@@ -38,33 +38,37 @@
  */
 public class StandbyTask extends AbstractTask implements Task {
     private final TaskId id;
+    private final Logger log;
+    private final String logPrefix;
     private final ProcessorTopology topology;
+    private final Set<TopicPartition> partitions;
     private final ProcessorStateManager stateMgr;
-    private final String logPrefix;
-    private final Logger log;
     private final StateDirectory stateDirectory;
     private final Sensor closeTaskSensor;
     private final InternalProcessorContext processorContext;
 
     /**
      * @param id             the ID of this task
+     * @param partitions     input topic partitions, used for thread metadata only
      * @param topology       the instance of {@link ProcessorTopology}
      * @param config         the {@link StreamsConfig} specified by the user
      * @param metrics        the {@link StreamsMetrics} created by the thread
+     * @param stateMgr       the {@link ProcessorStateManager} for this task
      * @param stateDirectory the {@link StateDirectory} created by the thread
      */
     StandbyTask(final TaskId id,
+                final Set<TopicPartition> partitions,
                 final ProcessorTopology topology,
                 final StreamsConfig config,
                 final StreamsMetricsImpl metrics,
                 final ProcessorStateManager stateMgr,
                 final StateDirectory stateDirectory) {
         this.id = id;
+        this.stateMgr = stateMgr;
         this.topology = topology;
+        this.partitions = partitions;
         this.stateDirectory = stateDirectory;
 
-        this.stateMgr = stateMgr;
-
         final String threadIdPrefix = String.format("stream-thread [%s] ", Thread.currentThread().getName());
         logPrefix = threadIdPrefix + String.format("%s [%s] ", "standby-task", id);
         final LogContext logContext = new LogContext(logPrefix);
@@ -194,7 +198,7 @@ public TaskId id() {
 
     @Override
     public Set<TopicPartition> inputPartitions() {
-        return Collections.emptySet();
+        return partitions;
     }
 
     @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StoreChangelogReader.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StoreChangelogReader.java
index 960d0b1da7fb5..ef2db354360e3 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StoreChangelogReader.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StoreChangelogReader.java
@@ -643,7 +643,7 @@ private void addChangelogsToRestoreConsumer(final Set<TopicPartition> partitions
 
         // the current assignment should not contain any of the new partitions
         if (assignment.removeAll(partitions)) {
-            throw new IllegalStateException("The current assignment " + assignment + " " +
+            throw new IllegalStateException("The current assignment " + restoreConsumer.assignment() + " " +
                 "already contains some of the new partitions " + partitions);
         }
         assignment.addAll(partitions);
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java
index 052827f780625..110077267508f 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java
@@ -20,7 +20,6 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.InvalidOffsetException;
 import org.apache.kafka.clients.producer.Producer;
@@ -53,7 +52,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -472,6 +470,7 @@ StandbyTask createTask(final Consumer<byte[], byte[]> consumer,
 
                 return new StandbyTask(
                     taskId,
+                    partitions,
                     topology,
                     config,
                     streamsMetrics,
@@ -513,7 +512,6 @@ StandbyTask createTask(final Consumer<byte[], byte[]> consumer,
     private volatile State state = State.CREATED;
     private volatile ThreadMetadata threadMetadata;
     private StreamThread.StateListener stateListener;
-    private Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> standbyRecords;
 
     private final ChangelogReader changelogReader;
 
@@ -633,11 +631,9 @@ public StreamThread(final Time time,
                         final LogContext logContext,
                         final AtomicInteger assignmentErrorCode) {
         super(threadId);
-        this.adminClient = adminClient;
-
         this.stateLock = new Object();
-        this.standbyRecords = new HashMap<>();
 
+        this.adminClient = adminClient;
         this.streamsMetrics = streamsMetrics;
         this.commitSensor = ThreadMetrics.commitSensor(threadId, streamsMetrics);
         this.pollSensor = ThreadMetrics.pollSensor(threadId, streamsMetrics);
@@ -1111,12 +1107,6 @@ private void completeShutdown(final boolean cleanRun) {
         log.info("Shutdown complete");
     }
 
-    void clearStandbyRecords(final List<TopicPartition> partitions) {
-        for (final TopicPartition tp : partitions) {
-            standbyRecords.remove(tp);
-        }
-    }
-
     /**
      * Return information about the current {@link StreamThread}.
      *
@@ -1167,7 +1157,7 @@ private void updateThreadMetadata(final Map<TaskId, Task> activeTasks,
             standbyTasksMetadata);
     }
 
-    public Iterable<Task> activeTasks() {
+    public List<Task> activeTasks() {
         return taskManager.activeTaskIterable();
     }
 
@@ -1238,15 +1228,11 @@ TaskManager taskManager() {
         return taskManager;
     }
 
-    Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> standbyRecords() {
-        return standbyRecords;
-    }
-
     int currentNumIterations() {
         return numIterations;
     }
 
-    public StreamThread.StateListener stateListener() {
-        return stateListener;
+    Throwable rebalanceException() {
+        return rebalanceException;
     }
 }
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
index 7b559a41bb4c0..8c3659566f8b3 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
@@ -265,7 +265,7 @@ protected static Set<TaskId> prepareForSubscription(final TaskManager taskManage
                 throw new IllegalStateException("Streams partition assignor's rebalance protocol is unknown");
         }
 
-        taskManager.handleRebalanceStart();
+        taskManager.handleRebalanceStart(topics);
 
         return activeTasks;
     }
@@ -1128,7 +1128,7 @@ public void onAssignment(final Assignment assignment, final ConsumerGroupMetadat
         // version 1 field
         final Map<TaskId, Set<TopicPartition>> activeTasks;
         // version 2 fields
-        final Map<TopicPartition, PartitionInfo> topicToPartitionInfo = new HashMap<>();
+        final Map<TopicPartition, PartitionInfo> topicToPartitionInfo;
         final Map<HostInfo, Set<TopicPartition>> partitionsByHost;
         final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost;
 
@@ -1139,6 +1139,7 @@ public void onAssignment(final Assignment assignment, final ConsumerGroupMetadat
                 activeTasks = getActiveTasks(partitions, info);
                 partitionsByHost = Collections.emptyMap();
                 standbyPartitionsByHost = Collections.emptyMap();
+                topicToPartitionInfo = Collections.emptyMap();
                 break;
             case 2:
             case 3:
@@ -1147,47 +1148,17 @@ public void onAssignment(final Assignment assignment, final ConsumerGroupMetadat
                 validateActiveTaskEncoding(partitions, info, logPrefix);
 
                 activeTasks = getActiveTasks(partitions, info);
-
-                // process partitions by host
-                for (final Set<TopicPartition> value : info.partitionsByHost().values()) {
-                    for (final TopicPartition topicPartition : value) {
-                        topicToPartitionInfo.put(
-                            topicPartition,
-                            new PartitionInfo(
-                                topicPartition.topic(),
-                                topicPartition.partition(),
-                                null,
-                                new Node[0],
-                                new Node[0]
-                            )
-                        );
-                    }
-                }
                 partitionsByHost = info.partitionsByHost();
                 standbyPartitionsByHost = Collections.emptyMap();
+                topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
                 break;
             case 6:
                 validateActiveTaskEncoding(partitions, info, logPrefix);
 
                 activeTasks = getActiveTasks(partitions, info);
-
-                // process partitions by host
-                for (final Set<TopicPartition> value : info.partitionsByHost().values()) {
-                    for (final TopicPartition topicPartition : value) {
-                        topicToPartitionInfo.put(
-                            topicPartition,
-                            new PartitionInfo(
-                                topicPartition.topic(),
-                                topicPartition.partition(),
-                                null,
-                                new Node[0],
-                                new Node[0]
-                            )
-                        );
-                    }
-                }
                 partitionsByHost = info.partitionsByHost();
                 standbyPartitionsByHost = info.standbyPartitionByHost();
+                topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
                 break;
             default:
                 throw new IllegalStateException(
@@ -1201,60 +1172,19 @@ public void onAssignment(final Assignment assignment, final ConsumerGroupMetadat
         taskManager.handleAssignment(activeTasks, info.standbyTasks());
     }
 
-    private Map<TaskId, Set<TopicPartition>> getActiveTasks(final List<TopicPartition> partitions, final AssignmentInfo info) {
-        final Map<TaskId, Set<TopicPartition>> activeTasks;
-        final Map<TaskId, Set<TopicPartition>> result =  new HashMap<>();
+    // protected for upgrade test
+    protected static Map<TaskId, Set<TopicPartition>> getActiveTasks(final List<TopicPartition> partitions, final AssignmentInfo info) {
+        final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
         for (int i = 0; i < partitions.size(); i++) {
             final TopicPartition partition = partitions.get(i);
             final TaskId id = info.activeTasks().get(i);
-            result.computeIfAbsent(id, k1 -> new HashSet<>()).add(partition);
+            activeTasks.computeIfAbsent(id, k1 -> new HashSet<>()).add(partition);
         }
-        activeTasks = result;
         return activeTasks;
     }
 
-    private static void validateActiveTaskEncoding(final List<TopicPartition> partitions, final AssignmentInfo info, final String logPrefix) {
-        // the number of assigned partitions should be the same as number of active tasks, which
-        // could be duplicated if one task has more than one assigned partitions
-        if (partitions.size() != info.activeTasks().size()) {
-            throw new TaskAssignmentException(
-                String.format(
-                    "%sNumber of assigned partitions %d is not equal to "
-                        + "the number of active taskIds %d, assignmentInfo=%s",
-                    logPrefix, partitions.size(),
-                    info.activeTasks().size(), info.toString()
-                )
-            );
-        }
-    }
-
-    private static void processVersionOneAssignment(final String logPrefix,
-                                                    final AssignmentInfo info,
-                                                    final List<TopicPartition> partitions,
-                                                    final Map<TaskId, Set<TopicPartition>> activeTasks,
-                                                    final Map<TopicPartition, TaskId> partitionsToTaskId) {
-        // the number of assigned partitions should be the same as number of active tasks, which
-        // could be duplicated if one task has more than one assigned partitions
-        validateActiveTaskEncoding(partitions, info, logPrefix);
-
-        for (int i = 0; i < partitions.size(); i++) {
-            final TopicPartition partition = partitions.get(i);
-            final TaskId id = info.activeTasks().get(i);
-            activeTasks.computeIfAbsent(id, k -> new HashSet<>()).add(partition);
-            partitionsToTaskId.put(partition, id);
-        }
-    }
-
-    public static void processVersionTwoAssignment(final String logPrefix,
-                                                   final AssignmentInfo info,
-                                                   final List<TopicPartition> partitions,
-                                                   final Map<TaskId, Set<TopicPartition>> activeTasks,
-                                                   final Map<TopicPartition, PartitionInfo> topicToPartitionInfo,
-                                                   final Map<TopicPartition, TaskId> partitionsToTaskId) {
-        processVersionOneAssignment(logPrefix, info, partitions, activeTasks, partitionsToTaskId);
-
-        // process partitions by host
-        final Map<HostInfo, Set<TopicPartition>> partitionsByHost = info.partitionsByHost();
+    private static Map<TopicPartition, PartitionInfo> getTopicPartitionInfo(final Map<HostInfo, Set<TopicPartition>> partitionsByHost) {
+        final Map<TopicPartition, PartitionInfo> topicToPartitionInfo = new HashMap<>();
         for (final Set<TopicPartition> value : partitionsByHost.values()) {
             for (final TopicPartition topicPartition : value) {
                 topicToPartitionInfo.put(
@@ -1269,6 +1199,22 @@ public static void processVersionTwoAssignment(final String logPrefix,
                 );
             }
         }
+        return topicToPartitionInfo;
+    }
+
+    private static void validateActiveTaskEncoding(final List<TopicPartition> partitions, final AssignmentInfo info, final String logPrefix) {
+        // the number of assigned partitions should be the same as number of active tasks, which
+        // could be duplicated if one task has more than one assigned partitions
+        if (partitions.size() != info.activeTasks().size()) {
+            throw new TaskAssignmentException(
+                String.format(
+                    "%sNumber of assigned partitions %d is not equal to "
+                        + "the number of active taskIds %d, assignmentInfo=%s",
+                    logPrefix, partitions.size(),
+                    info.activeTasks().size(), info.toString()
+                )
+            );
+        }
     }
 
     /**
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java
index 5df14458c8d1e..ab3c54127802e 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java
@@ -96,7 +96,7 @@ public void onPartitionsLost(final Collection<TopicPartition> partitions) {
         final long start = time.milliseconds();
         try {
             // close all active tasks as lost but don't try to commit offsets as we no longer own them
-            taskManager.handleTaskLoss();
+            taskManager.handleLostAll();
         } catch (final Throwable t) {
             log.error(
                 "Error caught during partitions lost, " +
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java
index 3630dbf71f2b2..6e22263530bcb 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java
@@ -32,6 +32,37 @@ public interface Task {
 
     long LATEST_OFFSET = -2L;
 
+    /*
+     *
+     * <pre>
+     *                 +-------------+
+     *          +<---- | Created (0) |
+     *          |      +-----+-------+
+     *          |            |
+     *          |            v
+     *          |      +-----+-------+
+     *          +<---- | Restoring(1)|<---------------+
+     *          |      +-----+-------+                |
+     *          |            |                        |
+     *          |            +--------------------+   |
+     *          |            |                    |   |
+     *          |            v                    v   |
+     *          |      +-----+-------+       +----+---+----+
+     *          |      | Running (2) | ----> | Suspended(3)|   * //TODO Suspended(3) could be removed after we've stable on KIP-429
+     *          |      +-----+-------+       +------+------+
+     *          |            |                      |
+     *          |            |                      |
+     *          |            v                      |
+     *          |      +-----+-------+              |
+     *          +----> | Closing (4) | <------------+
+     *                 +-----+-------+
+     *                       |
+     *                       v
+     *                 +-----+-------+
+     *                 | Closed (5)  |
+     *                 +-------------+
+     * </pre>
+     */
     enum State {
         CREATED(1, 4) {         // 0
             @Override
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java
index 65c7ac2ba2f6e..fb64fba0eb72c 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java
@@ -101,7 +101,9 @@ InternalTopologyBuilder builder() {
         return builder;
     }
 
-    void handleRebalanceStart() {
+    void handleRebalanceStart(final Set<String> subscribedTopics) {
+        builder.addSubscribedTopicsFromMetadata(subscribedTopics, logPrefix);
+
         rebalanceInProgress = true;
     }
 
@@ -113,8 +115,14 @@ void handleRebalanceComplete() {
         rebalanceInProgress = false;
     }
 
-    void handleAssignment(final Map<TaskId, Set<TopicPartition>> activeTasks,
-                          final Map<TaskId, Set<TopicPartition>> standbyTasks) {
+    /**
+     * @throws TaskMigratedException if the task producer got fenced (EOS only)
+     * @throws StreamsException fatal error while creating / initializing the task
+     *
+     * public for upgrade testing only
+     */
+    public void handleAssignment(final Map<TaskId, Set<TopicPartition>> activeTasks,
+                                 final Map<TaskId, Set<TopicPartition>> standbyTasks) {
         log.info("Handle new assignment with:\n\tNew active tasks: {}\n\tNew standby tasks: {}" +
                 "\n\tExisting active tasks: {}\n\tExisting standby tasks: {}",
             activeTasks.keySet(), standbyTasks.keySet(), activeTaskIds(), standbyTaskIds());
@@ -241,7 +249,6 @@ void handleRevocation(final Collection<TopicPartition> revokedPartitions) {
             if (remainingPartitions.containsAll(task.inputPartitions())) {
                 revokedTasks.add(task.id());
                 remainingPartitions.removeAll(task.inputPartitions());
-                break;
             }
         }
 
@@ -252,11 +259,7 @@ void handleRevocation(final Collection<TopicPartition> revokedPartitions) {
 
         for (final TaskId taskId : revokedTasks) {
             final Task task = tasks.get(taskId);
-            try {
-                task.suspend();
-            } catch (final RuntimeException e) {
-                throw new StreamsException("Unexpected exception while suspending task " + taskId, e);
-            }
+            task.suspend();
         }
     }
 
@@ -265,7 +268,7 @@ void handleRevocation(final Collection<TopicPartition> revokedPartitions) {
      * NOTE this method assumes that when it is called, EVERY task/partition has been lost and must
      * be closed as a zombie.
      */
-    void handleTaskLoss() {
+    void handleLostAll() {
         log.debug("Closing lost active tasks as zombies.");
 
         final Iterator<Task> iterator = tasks.values().iterator();
@@ -369,15 +372,6 @@ Task taskForInputPartition(final TopicPartition partition) {
         return partitionToTask.get(partition);
     }
 
-    StandbyTask standbyTask(final TopicPartition partition) {
-        for (final Task task : (Iterable<Task>) standbyTaskStream()::iterator) {
-            if (task.inputPartitions().contains(partition)) {
-                return (StandbyTask) task;
-            }
-        }
-        return null;
-    }
-
     Map<TaskId, Task> tasks() {
         // not bothering with an unmodifiable map, since the tasks themselves are mutable, but
         // if any outside code modifies the map or the tasks, it would be a severe transgression.
@@ -388,8 +382,8 @@ Map<TaskId, Task> activeTaskMap() {
         return activeTaskStream().collect(Collectors.toMap(Task::id, t -> t));
     }
 
-    Iterable<Task> activeTaskIterable() {
-        return activeTaskStream()::iterator;
+    List<Task> activeTaskIterable() {
+        return activeTaskStream().collect(Collectors.toList());
     }
 
     private Stream<Task> activeTaskStream() {
@@ -546,14 +540,19 @@ public String toString(final String indent) {
         return stringBuilder.toString();
     }
 
-    // FIXME: this is used from StreamThread only for a hack to collect metrics from the record collectors inside of StreamTasks
+    // below are for testing only
+    StandbyTask standbyTask(final TopicPartition partition) {
+        for (final Task task : (Iterable<Task>) standbyTaskStream()::iterator) {
+            if (task.inputPartitions().contains(partition)) {
+                return (StandbyTask) task;
+            }
+        }
+        return null;
+    }
+
+    // TODO: this is used from StreamThread only for a hack to collect metrics from the record collectors inside of StreamTasks
     // Instead, we should register and record the metrics properly inside of the record collector.
     Map<TaskId, StreamTask> fixmeStreamTasks() {
         return tasks.values().stream().filter(t -> t instanceof StreamTask).map(t -> (StreamTask) t).collect(Collectors.toMap(Task::id, t -> t));
     }
-
-    // FIXME: inappropriately used from StreamsUpgradeTest
-    public void fixmeUpdateSubscriptionsFromAssignment(final List<TopicPartition> partitions) {
-        builder.addSubscribedTopicsFromAssignment(partitions, logPrefix);
-    }
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java
index 324bb82d7e77a..8a761641746c0 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/GlobalThreadShutDownOrderTest.java
@@ -17,11 +17,11 @@
 
 package org.apache.kafka.streams.integration;
 
-import kafka.utils.MockTime;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.common.serialization.LongSerializer;
 import org.apache.kafka.common.serialization.Serdes;
 import org.apache.kafka.common.serialization.StringSerializer;
+import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.streams.KafkaStreams;
 import org.apache.kafka.streams.KeyValue;
@@ -38,18 +38,22 @@
 import org.apache.kafka.test.IntegrationTest;
 import org.apache.kafka.test.MockProcessorSupplier;
 import org.apache.kafka.test.TestUtils;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import static org.junit.Assert.assertEquals;
+
 
 /**
  * This test asserts that when Kafka Streams is closing and shuts
@@ -65,8 +69,6 @@ public class GlobalThreadShutDownOrderTest {
 
     private static final int NUM_BROKERS = 1;
     private static final Properties BROKER_CONFIG;
-    private final AtomicInteger closeCounter = new AtomicInteger(0);
-    private final int expectedCloseCount = 1;
 
     static {
         BROKER_CONFIG = new Properties();
@@ -74,6 +76,8 @@ public class GlobalThreadShutDownOrderTest {
         BROKER_CONFIG.put("transaction.state.log.min.isr", 1);
     }
 
+    private final AtomicInteger closeCounter = new AtomicInteger(0);
+
     @ClassRule
     public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(NUM_BROKERS, BROKER_CONFIG);
 
@@ -121,32 +125,31 @@ public void before() throws Exception {
     }
 
     @After
-    public void whenShuttingDown() throws Exception {
+    public void after() throws Exception {
         if (kafkaStreams != null) {
             kafkaStreams.close();
         }
         IntegrationTestUtils.purgeLocalStreamsState(streamsConfiguration);
     }
 
-    // FIXME
     @Test
     public void shouldFinishGlobalStoreOperationOnShutDown() throws Exception {
-//        kafkaStreams = new KafkaStreams(builder.build(), streamsConfiguration);
-//        populateTopics(globalStoreTopic);
-//        populateTopics(streamTopic);
-//
-//        kafkaStreams.start();
-//
-//        TestUtils.waitForCondition(
-//            () -> firstRecordProcessed,
-//            30000,
-//            "Has not processed record within 30 seconds");
-//
-//        kafkaStreams.close(Duration.ofSeconds(30));
-//
-//        final List<Long> expectedRetrievedValues = Arrays.asList(1L, 2L, 3L, 4L);
-//        assertEquals(expectedRetrievedValues, retrievedValuesList);
-//        assertEquals(expectedCloseCount, closeCounter.get());
+        kafkaStreams = new KafkaStreams(builder.build(), streamsConfiguration);
+        populateTopics(globalStoreTopic);
+        populateTopics(streamTopic);
+
+        kafkaStreams.start();
+
+        TestUtils.waitForCondition(
+            () -> firstRecordProcessed,
+            30000,
+            "Has not processed record within 30 seconds");
+
+        kafkaStreams.close(Duration.ofSeconds(30));
+
+        final List<Long> expectedRetrievedValues = Arrays.asList(1L, 2L, 3L, 4L);
+        assertEquals(expectedRetrievedValues, retrievedValuesList);
+        assertEquals(1, closeCounter.get());
     }
 
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/OptimizedKTableIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/OptimizedKTableIntegrationTest.java
index e2ef3f13f1531..570a5f12a1efa 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/OptimizedKTableIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/OptimizedKTableIntegrationTest.java
@@ -18,9 +18,7 @@
 
 import static org.apache.kafka.streams.integration.utils.IntegrationTestUtils.startApplicationAndWaitUntilRunning;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
 
@@ -31,7 +29,6 @@
 import java.util.Properties;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -87,39 +84,6 @@ public void after() {
         }
     }
 
-    @Test
-    public void standbyShouldNotPerformRestoreAtStartup() throws Exception {
-        final int numMessages = 10;
-        final int key = 1;
-        final Semaphore semaphore = new Semaphore(0);
-
-        final StreamsBuilder builder = new StreamsBuilder();
-        builder
-            .table(INPUT_TOPIC_NAME, Consumed.with(Serdes.Integer(), Serdes.Integer()),
-                Materialized.<Integer, Integer, KeyValueStore<Bytes, byte[]>>as(TABLE_NAME)
-                    .withCachingDisabled())
-            .toStream()
-            .peek((k, v) -> semaphore.release());
-
-        final KafkaStreams kafkaStreams1 = createKafkaStreams(builder, streamsConfiguration());
-        final KafkaStreams kafkaStreams2 = createKafkaStreams(builder, streamsConfiguration());
-        final List<KafkaStreams> kafkaStreamsList = Arrays.asList(kafkaStreams1, kafkaStreams2);
-
-        produceValueRange(key, 0, 10);
-
-        final AtomicLong restoreStartOffset = new AtomicLong(-1);
-        kafkaStreamsList.forEach(kafkaStreams -> {
-            kafkaStreams.setGlobalStateRestoreListener(createTrackingRestoreListener(restoreStartOffset, new AtomicLong()));
-        });
-        startApplicationAndWaitUntilRunning(kafkaStreamsList, Duration.ofSeconds(60));
-
-        // Assert that all messages in the first batch were processed in a timely manner
-        assertThat(semaphore.tryAcquire(numMessages, 60, TimeUnit.SECONDS), is(equalTo(true)));
-
-        // Assert that no restore occurred
-        assertThat(restoreStartOffset.get(), is(equalTo(-1L)));
-    }
-
     @Test
     public void shouldApplyUpdatesToStandbyStore() throws Exception {
         final int batch1NumMessages = 100;
@@ -138,11 +102,10 @@ public void shouldApplyUpdatesToStandbyStore() throws Exception {
         final KafkaStreams kafkaStreams1 = createKafkaStreams(builder, streamsConfiguration());
         final KafkaStreams kafkaStreams2 = createKafkaStreams(builder, streamsConfiguration());
         final List<KafkaStreams> kafkaStreamsList = Arrays.asList(kafkaStreams1, kafkaStreams2);
+        final TrackingStateRestoreListener listener = new TrackingStateRestoreListener();
 
-        final AtomicLong restoreStartOffset = new AtomicLong(-1L);
-        final AtomicLong restoreEndOffset = new AtomicLong(-1L);
         kafkaStreamsList.forEach(kafkaStreams -> {
-            kafkaStreams.setGlobalStateRestoreListener(createTrackingRestoreListener(restoreStartOffset, restoreEndOffset));
+            kafkaStreams.setGlobalStateRestoreListener(listener);
         });
         startApplicationAndWaitUntilRunning(kafkaStreamsList, Duration.ofSeconds(60));
 
@@ -170,7 +133,8 @@ public void shouldApplyUpdatesToStandbyStore() throws Exception {
 
         // Assert that no restore has occurred, ensures that when we check later that the restore
         // notification actually came from after the rebalance.
-        assertThat(restoreStartOffset.get(), is(equalTo(-1L)));
+        assertThat(listener.startOffset, is(equalTo(0L)));
+        assertThat(listener.totalNumRestored, is(equalTo(0L)));
 
         // Assert that the current value in store reflects all messages being processed
         assertThat(kafkaStreams1WasFirstActive ? store1.get(key) : store2.get(key), is(equalTo(batch1NumMessages - 1)));
@@ -195,11 +159,8 @@ public void shouldApplyUpdatesToStandbyStore() throws Exception {
         // Assert that all messages in the second batch were processed in a timely manner
         assertThat(semaphore.tryAcquire(batch2NumMessages, 60, TimeUnit.SECONDS), is(equalTo(true)));
 
-        // Assert that either restore was unnecessary or we restored from an offset later than 0
-        assertThat(restoreStartOffset.get(), is(anyOf(greaterThan(0L), equalTo(-1L))));
-
         // Assert that either restore was unnecessary or we restored to the last offset before we closed the kafkaStreams
-        assertThat(restoreEndOffset.get(), is(anyOf(equalTo(batch1NumMessages - 1L), equalTo(-1L))));
+        assertThat(listener.totalNumRestored, is((long) batch1NumMessages));
 
         // Assert that the current value in store reflects all messages being processed
         assertThat(newActiveStore.get(key), is(equalTo(totalNumMessages - 1)));
@@ -226,30 +187,32 @@ private KafkaStreams createKafkaStreams(final StreamsBuilder builder, final Prop
         return streams;
     }
 
-    private StateRestoreListener createTrackingRestoreListener(final AtomicLong restoreStartOffset,
-                                                               final AtomicLong restoreEndOffset) {
-        return new StateRestoreListener() {
-            @Override
-            public void onRestoreStart(final TopicPartition topicPartition,
-                                       final String storeName,
-                                       final long startingOffset,
-                                       final long endingOffset) {
-                restoreStartOffset.set(startingOffset);
-                restoreEndOffset.set(endingOffset);
-            }
-
-            @Override
-            public void onBatchRestored(final TopicPartition topicPartition, final String storeName,
-                final long batchEndOffset, final long numRestored) {
-
-            }
-
-            @Override
-            public void onRestoreEnd(final TopicPartition topicPartition, final String storeName,
-                final long totalRestored) {
-
-            }
-        };
+    private class TrackingStateRestoreListener implements StateRestoreListener {
+        long startOffset = -1L;
+        long endOffset = -1L;
+        long totalNumRestored = 0L;
+
+        @Override
+        public void onRestoreStart(final TopicPartition topicPartition,
+                                   final String storeName,
+                                   final long startingOffset,
+                                   final long endingOffset) {
+            startOffset = startingOffset;
+            endOffset = endingOffset;
+        }
+
+        @Override
+        public void onBatchRestored(final TopicPartition topicPartition,
+                                    final String storeName,
+                                    final long batchEndOffset,
+                                    final long numRestored) {
+            totalNumRestored += numRestored;
+        }
+
+        @Override
+        public void onRestoreEnd(final TopicPartition topicPartition, final String storeName, final long totalRestored) {
+
+        }
     }
 
     private Properties streamsConfiguration() {
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java
index 2d88ff3dc7a47..a818c12b64401 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/RestoreIntegrationTest.java
@@ -129,10 +129,11 @@ public void shouldRestoreStateFromSourceTopic() throws Exception {
         setCommittedOffset(INPUT_STREAM, offsetLimitDelta);
 
         final StateDirectory stateDirectory = new StateDirectory(new StreamsConfig(props), new MockTime(), true);
+        // note here the checkpointed offset is the last processed record's offset, so without control message we should write this offset - 1
         new OffsetCheckpoint(new File(stateDirectory.directoryForTask(new TaskId(0, 0)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(INPUT_STREAM, 0), (long) offsetCheckpointed));
+                .write(Collections.singletonMap(new TopicPartition(INPUT_STREAM, 0), (long) offsetCheckpointed - 1));
         new OffsetCheckpoint(new File(stateDirectory.directoryForTask(new TaskId(0, 1)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(INPUT_STREAM, 1), (long) offsetCheckpointed));
+                .write(Collections.singletonMap(new TopicPartition(INPUT_STREAM, 1), (long) offsetCheckpointed - 1));
 
         final CountDownLatch startupLatch = new CountDownLatch(1);
         final CountDownLatch shutdownLatch = new CountDownLatch(1);
@@ -191,10 +192,11 @@ public void shouldRestoreStateFromChangelogTopic() throws Exception {
         createStateForRestoration(INPUT_STREAM);
 
         final StateDirectory stateDirectory = new StateDirectory(new StreamsConfig(props), new MockTime(), true);
+        // note here the checkpointed offset is the last processed record's offset, so without control message we should write this offset - 1
         new OffsetCheckpoint(new File(stateDirectory.directoryForTask(new TaskId(0, 0)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(APPID + "-store-changelog", 0), (long) offsetCheckpointed));
+                .write(Collections.singletonMap(new TopicPartition(APPID + "-store-changelog", 0), (long) offsetCheckpointed - 1));
         new OffsetCheckpoint(new File(stateDirectory.directoryForTask(new TaskId(0, 1)), ".checkpoint"))
-                .write(Collections.singletonMap(new TopicPartition(APPID + "-store-changelog", 1), (long) offsetCheckpointed));
+                .write(Collections.singletonMap(new TopicPartition(APPID + "-store-changelog", 1), (long) offsetCheckpointed - 1));
 
         final CountDownLatch startupLatch = new CountDownLatch(1);
         final CountDownLatch shutdownLatch = new CountDownLatch(1);
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/SmokeTestDriverIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/SmokeTestDriverIntegrationTest.java
index 822a004715ba2..c2bd7a5d24615 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/SmokeTestDriverIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/SmokeTestDriverIntegrationTest.java
@@ -16,15 +16,22 @@
  */
 package org.apache.kafka.streams.integration;
 
+import org.apache.kafka.streams.StreamsConfig;
 import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
+import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
+import org.apache.kafka.streams.tests.SmokeTestClient;
 import org.apache.kafka.streams.tests.SmokeTestDriver;
 import org.apache.kafka.test.IntegrationTest;
+
+import org.junit.Assert;
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Map;
+import java.util.Properties;
 import java.util.Set;
 
 import static org.apache.kafka.streams.tests.SmokeTestDriver.generate;
@@ -71,71 +78,69 @@ SmokeTestDriver.VerificationResult result() {
 
     }
 
-    // FIXME
-
     @Test
     public void shouldWorkWithRebalance() throws InterruptedException {
-//        int numClientsCreated = 0;
-//        final ArrayList<SmokeTestClient> clients = new ArrayList<>();
-//
-//        IntegrationTestUtils.cleanStateBeforeTest(CLUSTER, SmokeTestDriver.topics());
-//
-//        final String bootstrapServers = CLUSTER.bootstrapServers();
-//        final Driver driver = new Driver(bootstrapServers, 10, 1000);
-//        driver.start();
-//        System.out.println("started driver");
-//
-//
-//        final Properties props = new Properties();
-//        props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
-//
-//        // cycle out Streams instances as long as the test is running.
-//        while (driver.isAlive()) {
-//            // take a nap
-//            Thread.sleep(1000);
-//
-//            // add a new client
-//            final SmokeTestClient smokeTestClient = new SmokeTestClient("streams-" + numClientsCreated++);
-//            clients.add(smokeTestClient);
-//            smokeTestClient.start(props);
-//
-//            while (!clients.get(clients.size() - 1).started()) {
-//                Thread.sleep(100);
-//            }
-//
-//            // let the oldest client die of "natural causes"
-//            if (clients.size() >= 3) {
-//                final SmokeTestClient client = clients.remove(0);
-//
-//                client.closeAsync();
-//                while (!client.closed()) {
-//                    Thread.sleep(100);
-//                }
-//            }
-//        }
-//
-//        try {
-//            // wait for verification to finish
-//            driver.join();
-//        } finally {
-//            // whether or not the assertions failed, tell all the streams instances to stop
-//            for (final SmokeTestClient client : clients) {
-//                client.closeAsync();
-//            }
-//
-//            // then, wait for them to stop
-//            for (final SmokeTestClient client : clients) {
-//                while (!client.closed()) {
-//                    Thread.sleep(100);
-//                }
-//            }
-//        }
-//
-//        // check to make sure that it actually succeeded
-//        if (driver.exception() != null) {
-//            driver.exception().printStackTrace();
-//            throw new AssertionError(driver.exception());
-//        }
-//        Assert.assertTrue(driver.result().result(), driver.result().passed());
+        int numClientsCreated = 0;
+        final ArrayList<SmokeTestClient> clients = new ArrayList<>();
+
+        IntegrationTestUtils.cleanStateBeforeTest(CLUSTER, SmokeTestDriver.topics());
+
+        final String bootstrapServers = CLUSTER.bootstrapServers();
+        final Driver driver = new Driver(bootstrapServers, 10, 1000);
+        driver.start();
+        System.out.println("started driver");
+
+
+        final Properties props = new Properties();
+        props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
+
+        // cycle out Streams instances as long as the test is running.
+        while (driver.isAlive()) {
+            // take a nap
+            Thread.sleep(1000);
+
+            // add a new client
+            final SmokeTestClient smokeTestClient = new SmokeTestClient("streams-" + numClientsCreated++);
+            clients.add(smokeTestClient);
+            smokeTestClient.start(props);
+
+            while (!clients.get(clients.size() - 1).started()) {
+                Thread.sleep(100);
+            }
+
+            // let the oldest client die of "natural causes"
+            if (clients.size() >= 3) {
+                final SmokeTestClient client = clients.remove(0);
+
+                client.closeAsync();
+                while (!client.closed()) {
+                    Thread.sleep(100);
+                }
+            }
+        }
+
+        try {
+            // wait for verification to finish
+            driver.join();
+        } finally {
+            // whether or not the assertions failed, tell all the streams instances to stop
+            for (final SmokeTestClient client : clients) {
+                client.closeAsync();
+            }
+
+            // then, wait for them to stop
+            for (final SmokeTestClient client : clients) {
+                while (!client.closed()) {
+                    Thread.sleep(100);
+                }
+            }
+        }
+
+        // check to make sure that it actually succeeded
+        if (driver.exception() != null) {
+            driver.exception().printStackTrace();
+            throw new AssertionError(driver.exception());
+        }
+        Assert.assertTrue(driver.result().result(), driver.result().passed());
     }
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/StreamStreamJoinIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/StreamStreamJoinIntegrationTest.java
index 5d2b6488df312..9be27baf15fb2 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/StreamStreamJoinIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/StreamStreamJoinIntegrationTest.java
@@ -16,407 +16,429 @@
  */
 package org.apache.kafka.streams.integration;
 
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.errors.InvalidStateStoreException;
+import org.apache.kafka.streams.kstream.Consumed;
+import org.apache.kafka.streams.kstream.JoinWindows;
+import org.apache.kafka.streams.kstream.KStream;
+import org.apache.kafka.streams.kstream.StreamJoined;
+import org.apache.kafka.streams.state.QueryableStoreTypes;
+import org.apache.kafka.streams.test.TestRecord;
 import org.apache.kafka.test.IntegrationTest;
+import org.apache.kafka.test.MockMapper;
 
+import org.junit.Before;
+import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 
-// FIXME
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+
+import static java.time.Duration.ofMillis;
+import static java.time.Duration.ofSeconds;
+import static org.junit.Assert.assertThrows;
+
 /**
  * Tests all available joins of Kafka Streams DSL.
  */
 @Category({IntegrationTest.class})
 @RunWith(value = Parameterized.class)
 public class StreamStreamJoinIntegrationTest extends AbstractJoinIntegrationTest {
-//    private KStream<Long, String> leftStream;
-//    private KStream<Long, String> rightStream;
+    private KStream<Long, String> leftStream;
+    private KStream<Long, String> rightStream;
 
     public StreamStreamJoinIntegrationTest(final boolean cacheEnabled) {
         super(cacheEnabled);
     }
-//
-//    @Before
-//    public void prepareTopology() throws InterruptedException {
-//        super.prepareEnvironment();
-//
-//        appID = "stream-stream-join-integration-test";
-//
-//        builder = new StreamsBuilder();
-//        leftStream = builder.stream(INPUT_TOPIC_LEFT);
-//        rightStream = builder.stream(INPUT_TOPIC_RIGHT);
-//    }
-//
-//    @Test
-//    public void shouldNotAccessJoinStoresWhenGivingName() throws InterruptedException {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-no-store-access");
-//        final StreamsBuilder builder = new StreamsBuilder();
-//
-//        final KStream<String, Integer> left = builder.stream(INPUT_TOPIC_LEFT, Consumed.with(Serdes.String(), Serdes.Integer()));
-//        final KStream<String, Integer> right = builder.stream(INPUT_TOPIC_RIGHT, Consumed.with(Serdes.String(), Serdes.Integer()));
-//        final CountDownLatch latch = new CountDownLatch(1);
-//
-//        left.join(
-//            right,
-//            (value1, value2) -> value1 + value2,
-//            JoinWindows.of(ofMillis(100)),
-//            StreamJoined.with(Serdes.String(), Serdes.Integer(), Serdes.Integer()).withStoreName("join-store"));
-//
-//        try (final KafkaStreams kafkaStreams = new KafkaStreams(builder.build(), STREAMS_CONFIG)) {
-//            kafkaStreams.setStateListener((newState, oldState) -> {
-//                if (newState == State.RUNNING) {
-//                    latch.countDown();
-//                }
-//            });
-//
-//            kafkaStreams.start();
-//            latch.await();
-//            assertThrows(InvalidStateStoreException.class, () -> kafkaStreams.store("join-store", QueryableStoreTypes.keyValueStore()));
-//        }
-//    }
-//
-//    @Test
-//    public void testInner() throws Exception {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-inner");
-//
-//        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
-//            null,
-//            null,
-//            null,
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
-//            null,
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
-//        );
-//
-//        leftStream.join(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
-//
-//        runTestWithDriver(expectedResult);
-//    }
-//
-//    @Test
-//    public void testInnerRepartitioned() throws Exception {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-inner-repartitioned");
-//
-//        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
-//            null,
-//            null,
-//            null,
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
-//            null,
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
-//        );
-//
-//        leftStream.map(MockMapper.noOpKeyValueMapper())
-//                .join(rightStream.flatMap(MockMapper.noOpFlatKeyValueMapper())
-//                                 .selectKey(MockMapper.selectKeyKeyValueMapper()),
-//                       valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
-//
-//        runTestWithDriver(expectedResult);
-//    }
-//
-//    @Test
-//    public void testLeft() throws Exception {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-left");
-//
-//        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
-//            null,
-//            null,
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
-//            null,
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
-//        );
-//
-//        leftStream.leftJoin(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
-//
-//        runTestWithDriver(expectedResult);
-//    }
-//
-//    @Test
-//    public void testLeftRepartitioned() throws Exception {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-left-repartitioned");
-//
-//        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
-//            null,
-//            null,
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
-//            null,
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
-//        );
-//
-//        leftStream.map(MockMapper.noOpKeyValueMapper())
-//                .leftJoin(rightStream.flatMap(MockMapper.noOpFlatKeyValueMapper())
-//                                     .selectKey(MockMapper.selectKeyKeyValueMapper()),
-//                        valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
-//
-//        runTestWithDriver(expectedResult);
-//    }
-//
-//    @Test
-//    public void testOuter() throws Exception {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-outer");
-//
-//        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
-//            null,
-//            null,
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
-//            null,
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
-//        );
-//
-//        leftStream.outerJoin(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
-//
-//        runTestWithDriver(expectedResult);
-//    }
-//
-//    @Test
-//    public void testOuterRepartitioned() throws Exception {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-outer");
-//
-//        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
-//            null,
-//            null,
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
-//            null,
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
-//        );
-//
-//        leftStream.map(MockMapper.noOpKeyValueMapper())
-//                .outerJoin(rightStream.flatMap(MockMapper.noOpFlatKeyValueMapper())
-//                                .selectKey(MockMapper.selectKeyKeyValueMapper()),
-//                        valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
-//
-//        runTestWithDriver(expectedResult);
-//    }
-//
-//    @Test
-//    public void testMultiInner() throws Exception {
-//        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-multi-inner");
-//
-//        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
-//            null,
-//            null,
-//            null,
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a-a", null, 4L)),
-//            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a-a", null, 5L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-a", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-a", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-a-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-a-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-b", null, 6L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-b", null, 6L)),
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-b", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-a", null, 9L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-b", null, 9L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-a", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-b", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-a", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-b", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-a", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-b", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-a-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-a-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-c", null, 10L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-c", null, 10L)),
-//            null,
-//            null,
-//            null,
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-a", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-b", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-c", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-a", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-b", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-c", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-a", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-b", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-c", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-a-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-a-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-d", null, 14L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-d", null, 14L)),
-//            Arrays.asList(
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-d", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-d", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-d", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-a", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-b", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-c", null, 15L),
-//                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-d", null, 15L))
-//        );
-//
-//        leftStream.join(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10)))
-//                .join(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
-//
-//        runTestWithDriver(expectedResult);
-//    }
+
+    @Before
+    public void prepareTopology() throws InterruptedException {
+        super.prepareEnvironment();
+
+        appID = "stream-stream-join-integration-test";
+
+        builder = new StreamsBuilder();
+        leftStream = builder.stream(INPUT_TOPIC_LEFT);
+        rightStream = builder.stream(INPUT_TOPIC_RIGHT);
+    }
+
+    @Test
+    public void shouldNotAccessJoinStoresWhenGivingName() throws InterruptedException {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-no-store-access");
+        final StreamsBuilder builder = new StreamsBuilder();
+
+        final KStream<String, Integer> left = builder.stream(INPUT_TOPIC_LEFT, Consumed.with(Serdes.String(), Serdes.Integer()));
+        final KStream<String, Integer> right = builder.stream(INPUT_TOPIC_RIGHT, Consumed.with(Serdes.String(), Serdes.Integer()));
+        final CountDownLatch latch = new CountDownLatch(1);
+
+        left.join(
+            right,
+            (value1, value2) -> value1 + value2,
+            JoinWindows.of(ofMillis(100)),
+            StreamJoined.with(Serdes.String(), Serdes.Integer(), Serdes.Integer()).withStoreName("join-store"));
+
+        try (final KafkaStreams kafkaStreams = new KafkaStreams(builder.build(), STREAMS_CONFIG)) {
+            kafkaStreams.setStateListener((newState, oldState) -> {
+                if (newState == KafkaStreams.State.RUNNING) {
+                    latch.countDown();
+                }
+            });
+
+            kafkaStreams.start();
+            latch.await();
+            assertThrows(InvalidStateStoreException.class, () -> kafkaStreams.store("join-store", QueryableStoreTypes.keyValueStore()));
+        }
+    }
+
+    @Test
+    public void testInner() {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-inner");
+
+        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
+            null,
+            null,
+            null,
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
+            null,
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
+        );
+
+        leftStream.join(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
+
+        runTestWithDriver(expectedResult);
+    }
+
+    @Test
+    public void testInnerRepartitioned() {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-inner-repartitioned");
+
+        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
+            null,
+            null,
+            null,
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
+            null,
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
+        );
+
+        leftStream.map(MockMapper.noOpKeyValueMapper())
+                .join(rightStream.flatMap(MockMapper.noOpFlatKeyValueMapper())
+                                 .selectKey(MockMapper.selectKeyKeyValueMapper()),
+                       valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
+
+        runTestWithDriver(expectedResult);
+    }
+
+    @Test
+    public void testLeft() {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-left");
+
+        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
+            null,
+            null,
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
+            null,
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
+        );
+
+        leftStream.leftJoin(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
+
+        runTestWithDriver(expectedResult);
+    }
+
+    @Test
+    public void testLeftRepartitioned() {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-left-repartitioned");
+
+        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
+            null,
+            null,
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
+            null,
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
+        );
+
+        leftStream.map(MockMapper.noOpKeyValueMapper())
+                .leftJoin(rightStream.flatMap(MockMapper.noOpFlatKeyValueMapper())
+                                     .selectKey(MockMapper.selectKeyKeyValueMapper()),
+                        valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
+
+        runTestWithDriver(expectedResult);
+    }
+
+    @Test
+    public void testOuter() {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-outer");
+
+        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
+            null,
+            null,
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
+            null,
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
+        );
+
+        leftStream.outerJoin(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
+
+        runTestWithDriver(expectedResult);
+    }
+
+    @Test
+    public void testOuterRepartitioned() {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-outer");
+
+        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
+            null,
+            null,
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-null", null, 3L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a", null, 4L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a", null, 5L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b", null, 6L)),
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b", null, 9L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c", null, 10L)),
+            null,
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d", null, 14L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d", null, 15L))
+        );
+
+        leftStream.map(MockMapper.noOpKeyValueMapper())
+                .outerJoin(rightStream.flatMap(MockMapper.noOpFlatKeyValueMapper())
+                                .selectKey(MockMapper.selectKeyKeyValueMapper()),
+                        valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
+
+        runTestWithDriver(expectedResult);
+    }
+
+    @Test
+    public void testMultiInner() {
+        STREAMS_CONFIG.put(StreamsConfig.APPLICATION_ID_CONFIG, appID + "-multi-inner");
+
+        final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
+            null,
+            null,
+            null,
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "A-a-a", null, 4L)),
+            Collections.singletonList(new TestRecord<>(ANY_UNIQUE_KEY, "B-a-a", null, 5L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-a", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-a", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-a-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-a-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-b", null, 6L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-b", null, 6L)),
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-b", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-a", null, 9L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-b", null, 9L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-a", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-b", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-a", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-b", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-a", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-b", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-a-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-a-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-c", null, 10L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-c", null, 10L)),
+            null,
+            null,
+            null,
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-a", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-b", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-c", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-a", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-b", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-c", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-a", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-b", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-c", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-a-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-a-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-b-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-b-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-a-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-b-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-c-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-c-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-c-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "A-d-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "B-d-d", null, 14L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "C-d-d", null, 14L)),
+            Arrays.asList(
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-a-d", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-b-d", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-c-d", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-a", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-b", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-c", null, 15L),
+                new TestRecord<>(ANY_UNIQUE_KEY, "D-d-d", null, 15L))
+        );
+
+        leftStream.join(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10)))
+                .join(rightStream, valueJoiner, JoinWindows.of(ofSeconds(10))).to(OUTPUT_TOPIC);
+
+        runTestWithDriver(expectedResult);
+    }
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/TableTableJoinIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/TableTableJoinIntegrationTest.java
index 7edce8797427f..59f968a757706 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/TableTableJoinIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/TableTableJoinIntegrationTest.java
@@ -174,8 +174,8 @@ public void testInnerInner() throws Exception {
         if (cacheEnabled) {
             runTestWithDriver(expectedFinalMultiJoinResult, storeName);
         } else {
-            // FIXME: the duplicate below for all the multi-joins
-            //        are due to KAFKA-6443, should be updated once it is fixed.
+            // TODO K6443: the duplicate below for all the multi-joins are due to
+            //             KAFKA-6443, should be updated once it is fixed.
             final List<List<TestRecord<Long, String>>> expectedResult = Arrays.asList(
                 null,
                 null,
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
index d09337aae952d..82f0d1b7edb98 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopologyBuilderTest.java
@@ -690,7 +690,7 @@ public void shouldSetCorrectSourceNodesWithRegexUpdatedTopics() {
         updatedTopics.add("topic-3");
         updatedTopics.add("topic-A");
 
-        builder.updateSubscribedTopics(updatedTopics, null);
+        builder.addSubscribedTopicsFromMetadata(updatedTopics, null);
         builder.setApplicationId("test-id");
 
         final Map<Integer, InternalTopologyBuilder.TopicsInfo> topicGroups = builder.topicGroups();
@@ -772,7 +772,7 @@ public void shouldConnectRegexMatchedTopicsToStateStore() {
         updatedTopics.add("topic-3");
         updatedTopics.add("topic-A");
 
-        builder.updateSubscribedTopics(updatedTopics, "test-thread");
+        builder.addSubscribedTopicsFromMetadata(updatedTopics, "test-thread");
         builder.setApplicationId("test-app");
 
         final Map<String, List<String>> stateStoreAndTopics = builder.stateStoreNameToSourceTopics();
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java
index c72984b07ed85..616556480bc05 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java
@@ -352,7 +352,7 @@ public void shouldThrowIfClosingOnIllegalState() {
     }
 
     private StandbyTask createStandbyTask() {
-        return new StandbyTask(taskId, topology, config, streamsMetrics, stateManager, stateDirectory);
+        return new StandbyTask(taskId, Collections.singleton(partition), topology, config, streamsMetrics, stateManager, stateDirectory);
     }
 
     private MetricName setupCloseTaskMetric() {
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
index c8b43545eb625..f4842f19fcce1 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
@@ -84,16 +84,13 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.stream.Stream;
 
-import static java.util.Collections.singletonList;
 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.mkProperties;
@@ -113,7 +110,6 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-// TODO K9113: this test needs to be fixed
 public class StreamThreadTest {
 
     private final static String APPLICATION_ID = "stream-thread-test";
@@ -165,11 +161,72 @@ private Properties configProps(final boolean enableEoS) {
         ));
     }
 
-    @Ignore
-    @Test
-    public void testPartitionAssignmentChangeForSingleGroup() {
-        internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
+    private Cluster createCluster() {
+        final Node node = new Node(-1, "localhost", 8121);
+        return new Cluster(
+            "mockClusterId",
+            Collections.singletonList(node),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            node
+        );
+    }
+
+    private StreamThread createStreamThread(@SuppressWarnings("SameParameterValue") final String clientId,
+                                            final StreamsConfig config,
+                                            final boolean eosEnabled) {
+        if (eosEnabled) {
+            clientSupplier.setApplicationIdForProducer(APPLICATION_ID);
+        }
+
+        clientSupplier.setClusterForAdminClient(createCluster());
+
+        final StreamsMetricsImpl streamsMetrics = new StreamsMetricsImpl(
+            metrics,
+            APPLICATION_ID,
+            config.getString(StreamsConfig.BUILT_IN_METRICS_VERSION_CONFIG)
+        );
 
+        return StreamThread.create(
+            internalTopologyBuilder,
+            config,
+            clientSupplier,
+            clientSupplier.getAdmin(config.getAdminConfigs(clientId)),
+            PROCESS_ID,
+            clientId,
+            streamsMetrics,
+            mockTime,
+            streamsMetadataState,
+            0,
+            stateDirectory,
+            new MockStateRestoreListener(),
+            threadIdx
+        );
+    }
+
+    private static class StateListenerStub implements StreamThread.StateListener {
+        int numChanges = 0;
+        ThreadStateTransitionValidator oldState = null;
+        ThreadStateTransitionValidator newState = null;
+
+        @Override
+        public void onChange(final Thread thread,
+                             final ThreadStateTransitionValidator newState,
+                             final ThreadStateTransitionValidator oldState) {
+            ++numChanges;
+            if (this.newState != null) {
+                if (this.newState != oldState) {
+                    throw new RuntimeException("State mismatch " + oldState + " different from " + this.newState);
+                }
+            }
+            this.oldState = oldState;
+            this.newState = newState;
+        }
+    }
+
+    @Test
+    public void shouldChangeStateInRebalanceListener() {
         final StreamThread thread = createStreamThread(CLIENT_ID, config, false);
 
         final StateListenerStub stateListener = new StateListenerStub();
@@ -189,9 +246,7 @@ public void testPartitionAssignmentChangeForSingleGroup() {
         assertEquals(thread.state(), StreamThread.State.PARTITIONS_REVOKED);
 
         // assign single partition
-        assignedPartitions = singletonList(t1p1);
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(Collections.emptyMap(), Collections.emptyMap());
+        assignedPartitions = Collections.singletonList(t1p1);
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(assignedPartitions);
@@ -207,7 +262,7 @@ public void testPartitionAssignmentChangeForSingleGroup() {
     }
 
     @Test
-    public void testStateChangeStartClose() throws Exception {
+    public void shouldChangeStateAtStartClose() throws Exception {
         final StreamThread thread = createStreamThread(CLIENT_ID, config, false);
 
         final StateListenerStub stateListener = new StateListenerStub();
@@ -229,61 +284,17 @@ public void testStateChangeStartClose() throws Exception {
         assertEquals(thread.state(), StreamThread.State.DEAD);
     }
 
-    private Cluster createCluster() {
-        final Node node = new Node(-1, "localhost", 8121);
-        return new Cluster(
-            "mockClusterId",
-            singletonList(node),
-            Collections.emptySet(),
-            Collections.emptySet(),
-            Collections.emptySet(),
-            node
-        );
-    }
-
-    private StreamThread createStreamThread(@SuppressWarnings("SameParameterValue") final String clientId,
-                                            final StreamsConfig config,
-                                            final boolean eosEnabled) {
-        if (eosEnabled) {
-            clientSupplier.setApplicationIdForProducer(APPLICATION_ID);
-        }
-
-        clientSupplier.setClusterForAdminClient(createCluster());
-
-        final StreamsMetricsImpl streamsMetrics = new StreamsMetricsImpl(
-            metrics,
-            APPLICATION_ID,
-            config.getString(StreamsConfig.BUILT_IN_METRICS_VERSION_CONFIG)
-        );
-
-        return StreamThread.create(
-            internalTopologyBuilder,
-            config,
-            clientSupplier,
-            clientSupplier.getAdmin(config.getAdminConfigs(clientId)),
-            PROCESS_ID,
-            clientId,
-            streamsMetrics,
-            mockTime,
-            streamsMetadataState,
-            0,
-            stateDirectory,
-            new MockStateRestoreListener(),
-            threadIdx
-        );
-    }
-
     @Test
-    public void testMetricsCreatedAtStartupWithBuiltInMetricsVersionLatest() {
-        testMetricsCreatedAtStartup(StreamsConfig.METRICS_LATEST);
+    public void shouldCreateMetricsAtStartupWithBuiltInMetricsVersionLatest() {
+        shouldCreateMetricsAtStartup(StreamsConfig.METRICS_LATEST);
     }
 
     @Test
-    public void testMetricsCreatedAtStartupWithBuiltInMetricsVersion0100To24() {
-        testMetricsCreatedAtStartup(StreamsConfig.METRICS_0100_TO_24);
+    public void shouldCreateMetricsAtStartupWithBuiltInMetricsVersion0100To24() {
+        shouldCreateMetricsAtStartup(StreamsConfig.METRICS_0100_TO_24);
     }
 
-    private void testMetricsCreatedAtStartup(final String builtInMetricsVersion) {
+    private void shouldCreateMetricsAtStartup(final String builtInMetricsVersion) {
         final Properties props = configProps(false);
         props.setProperty(StreamsConfig.BUILT_IN_METRICS_VERSION_CONFIG, builtInMetricsVersion);
         final StreamsConfig config = new StreamsConfig(props);
@@ -433,7 +444,6 @@ public void shouldNotCommitBeforeTheCommitInterval() {
         EasyMock.verify(taskManager);
     }
 
-    @Ignore
     @Test
     public void shouldRespectNumIterationsInMainLoop() {
         final MockProcessor mockProcessor = new MockProcessor(PunctuationType.WALL_CLOCK_TIME, 10L);
@@ -455,13 +465,11 @@ public void shouldRespectNumIterationsInMainLoop() {
 
         final TaskId task1 = new TaskId(0, t1p1.partition());
         final Set<TopicPartition> assignedPartitions = Collections.singleton(t1p1);
-//FIXME
-        //        thread.taskManager().setAssignmentMetadata(
-//            Collections.singletonMap(
-//                task1,
-//                assignedPartitions),
-//            Collections.emptyMap());
-//        thread.taskManager().setPartitionsToTaskId(Collections.singletonMap(t1p1, task1));
+
+        thread.taskManager().handleAssignment(
+            Collections.singletonMap(task1, assignedPartitions),
+            Collections.emptyMap()
+        );
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(Collections.singleton(t1p1));
@@ -603,16 +611,6 @@ public void shouldCommitAfterTheCommitInterval() {
         EasyMock.verify(taskManager);
     }
 
-    private TaskManager mockTaskManagerCommit(final Consumer<byte[], byte[]> consumer,
-                                              final int numberOfCommits,
-                                              final int commits) {
-        final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);
-        EasyMock.expect(taskManager.commitAll()).andReturn(commits).times(numberOfCommits);
-        EasyMock.replay(taskManager, consumer);
-        return taskManager;
-    }
-
-    @Ignore
     @Test
     public void shouldInjectSharedProducerForAllTasksUsingClientSupplierOnCreateIfEosDisabled() {
         internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
@@ -632,8 +630,7 @@ public void shouldInjectSharedProducerForAllTasksUsingClientSupplierOnCreateIfEo
         activeTasks.put(task1, Collections.singleton(t1p1));
         activeTasks.put(task2, Collections.singleton(t1p2));
 
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(assignedPartitions);
@@ -652,7 +649,6 @@ public void shouldInjectSharedProducerForAllTasksUsingClientSupplierOnCreateIfEo
         assertSame(clientSupplier.restoreConsumer, thread.restoreConsumer);
     }
 
-    @Ignore
     @Test
     public void shouldInjectProducerPerTaskUsingClientSupplierOnCreateIfEosEnable() {
         internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
@@ -664,19 +660,14 @@ public void shouldInjectProducerPerTaskUsingClientSupplierOnCreateIfEosEnable()
 
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
         final List<TopicPartition> assignedPartitions = new ArrayList<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
 
         // assign single partition
         assignedPartitions.add(t1p1);
         assignedPartitions.add(t1p2);
         activeTasks.put(task1, Collections.singleton(t1p1));
         activeTasks.put(task2, Collections.singleton(t1p2));
-        partitionsToTaskId.put(t1p1, task1);
-        partitionsToTaskId.put(t1p2, task2);
 
-        // TODO K9113: fix this test
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(assignedPartitions);
@@ -688,19 +679,23 @@ public void shouldInjectProducerPerTaskUsingClientSupplierOnCreateIfEosEnable()
 
         thread.runOnce();
 
-        assertEquals(Stream.of(thread.activeTasks()).count(), clientSupplier.producers.size());
+        assertEquals(thread.activeTasks().size(), clientSupplier.producers.size());
         assertSame(clientSupplier.consumer, thread.consumer);
         assertSame(clientSupplier.restoreConsumer, thread.restoreConsumer);
     }
 
-    @Ignore
     @Test
-    public void shouldCloseAllTaskProducersOnCloseIfEosEnabled() {
+    public void shouldCloseAllTaskProducersOnCloseIfEosEnabled() throws InterruptedException {
         internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
 
         final StreamThread thread = createStreamThread(CLIENT_ID, new StreamsConfig(configProps(true)), true);
 
-        thread.setState(StreamThread.State.STARTING);
+        thread.start();
+        TestUtils.waitForCondition(
+            () -> thread.state() == StreamThread.State.STARTING,
+            10 * 1000,
+            "Thread never started.");
+
         thread.rebalanceListener.onPartitionsRevoked(Collections.emptyList());
 
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
@@ -712,19 +707,18 @@ public void shouldCloseAllTaskProducersOnCloseIfEosEnabled() {
         activeTasks.put(task1, Collections.singleton(t1p1));
         activeTasks.put(task2, Collections.singleton(t1p2));
 
-// TODO K9113: fix this test
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
-        final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
-        mockConsumer.assign(assignedPartitions);
-        final Map<TopicPartition, Long> beginOffsets = new HashMap<>();
-        beginOffsets.put(t1p1, 0L);
-        beginOffsets.put(t1p2, 0L);
-        mockConsumer.updateBeginningOffsets(beginOffsets);
-
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
         thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
 
+        for (final Task task : thread.activeTasks()) {
+            assertTrue(((MockProducer) ((RecordCollectorImpl) ((StreamTask) task).recordCollector()).producer()).transactionInitialized());
+        }
+
         thread.shutdown();
-        thread.run();
+        TestUtils.waitForCondition(
+            () -> thread.state() == StreamThread.State.DEAD,
+            10 * 1000,
+            "Thread never shut down.");
 
         for (final Task task : thread.activeTasks()) {
             assertTrue(((MockProducer) ((RecordCollectorImpl) ((StreamTask) task).recordCollector()).producer()).closed());
@@ -829,9 +823,8 @@ public void shouldOnlyShutdownOnce() {
         EasyMock.verify(taskManager);
     }
 
-    @Ignore
     @Test
-    public void shouldNotNullPointerWhenStandbyTasksAssignedAndNoStateStoresForTopology() {
+    public void shouldNotThrowWhenStandbyTasksAssignedAndNoStateStoresForTopology() {
         internalTopologyBuilder.addSource(null, "name", null, null, null, "topic");
         internalTopologyBuilder.addSink("out", "output", null, null, null, "name");
 
@@ -845,24 +838,13 @@ public void shouldNotNullPointerWhenStandbyTasksAssignedAndNoStateStoresForTopol
         // assign single partition
         standbyTasks.put(task1, Collections.singleton(t1p1));
 
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(Collections.emptyMap(), standbyTasks);
-//        thread.taskManager().createTasks(Collections.emptyList());
+        thread.taskManager().handleAssignment(Collections.emptyMap(), standbyTasks);
 
         thread.rebalanceListener.onPartitionsAssigned(Collections.emptyList());
     }
 
-    public List<Task> asList(final Iterable<Task> tasks) {
-        final List<Task> result = new LinkedList<>();
-        for (final Task task : tasks) {
-            result.add(task);
-        }
-        return result;
-    }
-
-    @Ignore
     @Test
-    public void shouldNotCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerWasFencedWhileProcessing() throws Exception {
+    public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerWasFencedWhileProcessing() throws Exception {
         internalTopologyBuilder.addSource(null, "source", null, null, null, topic1);
         internalTopologyBuilder.addSink("sink", "dummyTopic", null, null, null, "source");
 
@@ -870,23 +852,19 @@ public void shouldNotCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerWasFence
 
         final MockConsumer<byte[], byte[]> consumer = clientSupplier.consumer;
 
-        consumer.updatePartitions(topic1, singletonList(new PartitionInfo(topic1, 1, null, null, null)));
+        consumer.updatePartitions(topic1, Collections.singletonList(new PartitionInfo(topic1, 1, null, null, null)));
 
         thread.setState(StreamThread.State.STARTING);
         thread.rebalanceListener.onPartitionsRevoked(Collections.emptySet());
 
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
         final List<TopicPartition> assignedPartitions = new ArrayList<>();
 
         // assign single partition
         assignedPartitions.add(t1p1);
         activeTasks.put(task1, Collections.singleton(t1p1));
-        partitionsToTaskId.put(t1p1, task1);
 
-        // TODO K9113: fix this test
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(assignedPartitions);
@@ -894,8 +872,7 @@ public void shouldNotCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerWasFence
         thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
 
         thread.runOnce();
-        final Iterable<Task> tasks = thread.activeTasks();
-        assertThat(asList(tasks).size(), equalTo(1));
+        assertThat(thread.activeTasks().size(), equalTo(1));
         final MockProducer producer = clientSupplier.producers.get(0);
 
         // change consumer subscription from "pattern" to "manual" to be able to call .addRecords()
@@ -920,17 +897,16 @@ public void shouldNotCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerWasFence
             thread.runOnce();
             fail("Should have thrown TaskMigratedException");
         } catch (final KafkaException expected) {
-            assertTrue(expected.getCause() instanceof TaskMigratedException);
+            assertTrue(expected instanceof TaskMigratedException);
             assertTrue("StreamsThread removed the fenced zombie task already, should wait for rebalance to close all zombies together.",
-                       asList(thread.activeTasks()).stream().anyMatch(task -> task.id().equals(task1)));
+                thread.activeTasks().stream().anyMatch(task -> task.id().equals(task1)));
         }
 
         assertThat(producer.commitCount(), equalTo(1L));
     }
 
-    @Ignore
     @Test
-    public void shouldCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerGotFencedInCommitTransactionWhenSuspendingTasks() {
+    public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerGotFencedInCommitTransactionWhenSuspendingTasks() {
         final StreamThread thread = createStreamThread(CLIENT_ID, new StreamsConfig(configProps(true)), true);
 
         internalTopologyBuilder.addSource(null, "name", null, null, null, topic1);
@@ -940,17 +916,13 @@ public void shouldCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerGotFencedIn
         thread.rebalanceListener.onPartitionsRevoked(Collections.emptySet());
 
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
         final List<TopicPartition> assignedPartitions = new ArrayList<>();
 
         // assign single partition
         assignedPartitions.add(t1p1);
         activeTasks.put(task1, Collections.singleton(t1p1));
-        partitionsToTaskId.put(t1p1, task1);
 
-        // TODO K9113: fix this test
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(assignedPartitions);
@@ -959,39 +931,38 @@ public void shouldCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerGotFencedIn
 
         thread.runOnce();
 
-        assertThat(asList(thread.activeTasks()).size(), equalTo(1));
+        assertThat(thread.activeTasks().size(), equalTo(1));
 
-        clientSupplier.producers.get(0).fenceProducer();
+        clientSupplier.producers.get(0).fenceProducerOnCommitTxn();
         thread.rebalanceListener.onPartitionsRevoked(assignedPartitions);
+        assertTrue(thread.rebalanceException() instanceof TaskMigratedException);
         assertFalse(clientSupplier.producers.get(0).transactionCommitted());
-        assertTrue(clientSupplier.producers.get(0).closed());
-        assertTrue(asList(thread.activeTasks()).isEmpty());
-        assertTrue(asList(thread.activeTasks()).isEmpty());
+        assertFalse(clientSupplier.producers.get(0).closed());
+        assertEquals(1, thread.activeTasks().size());
     }
 
-    @Ignore
     @Test
-    public void shouldCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerGotFencedInCloseTransactionWhenSuspendingTasks() {
+    public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerGotFencedInCommitTransactionWhenCommitting() {
+        // only have source but no sink so that we would not get fenced in producer.send
+        internalTopologyBuilder.addSource(null, "source", null, null, null, topic1);
+
         final StreamThread thread = createStreamThread(CLIENT_ID, new StreamsConfig(configProps(true)), true);
 
-        internalTopologyBuilder.addSource(null, "name", null, null, null, topic1);
-        internalTopologyBuilder.addSink("out", "output", null, null, null, "name");
+        final MockConsumer<byte[], byte[]> consumer = clientSupplier.consumer;
+
+        consumer.updatePartitions(topic1, Collections.singletonList(new PartitionInfo(topic1, 1, null, null, null)));
 
         thread.setState(StreamThread.State.STARTING);
         thread.rebalanceListener.onPartitionsRevoked(Collections.emptySet());
 
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
         final List<TopicPartition> assignedPartitions = new ArrayList<>();
 
         // assign single partition
         assignedPartitions.add(t1p1);
         activeTasks.put(task1, Collections.singleton(t1p1));
-        partitionsToTaskId.put(t1p1, task1);
 
-// TODO K9113: fix this test
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(assignedPartitions);
@@ -999,39 +970,63 @@ public void shouldCloseTaskAsZombieAndRemoveFromActiveTasksIfProducerGotFencedIn
         thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
 
         thread.runOnce();
+        assertThat(thread.activeTasks().size(), equalTo(1));
+        final MockProducer producer = clientSupplier.producers.get(0);
 
-        assertThat(asList(thread.activeTasks()).size(), equalTo(1));
+        producer.fenceProducerOnCommitTxn();
+        mockTime.sleep(config.getLong(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG) + 1L);
+        consumer.addRecord(new ConsumerRecord<>(topic1, 1, 1, new byte[0], new byte[0]));
+        try {
+            thread.runOnce();
+            fail("Should have thrown TaskMigratedException");
+        } catch (final KafkaException expected) {
+            assertTrue(expected instanceof TaskMigratedException);
+            assertTrue("StreamsThread removed the fenced zombie task already, should wait for rebalance to close all zombies together.",
+                thread.activeTasks().stream().anyMatch(task -> task.id().equals(task1)));
+        }
 
-        clientSupplier.producers.get(0).fenceProducerOnClose();
-        thread.rebalanceListener.onPartitionsRevoked(assignedPartitions);
+        assertThat(producer.commitCount(), equalTo(0L));
 
-        assertFalse(clientSupplier.producers.get(0).transactionInFlight());
-        assertTrue(clientSupplier.producers.get(0).transactionCommitted());
+        assertTrue(clientSupplier.producers.get(0).transactionInFlight());
+        assertFalse(clientSupplier.producers.get(0).transactionCommitted());
         assertFalse(clientSupplier.producers.get(0).closed());
-        assertTrue(asList(thread.activeTasks()).isEmpty());
+        assertEquals(1, thread.activeTasks().size());
     }
 
-    private static class StateListenerStub implements StreamThread.StateListener {
-        int numChanges = 0;
-        ThreadStateTransitionValidator oldState = null;
-        ThreadStateTransitionValidator newState = null;
+    @Test
+    public void shouldNotCloseTaskProducerWhenSuspending() {
+        final StreamThread thread = createStreamThread(CLIENT_ID, new StreamsConfig(configProps(true)), true);
 
-        @Override
-        public void onChange(final Thread thread,
-                             final ThreadStateTransitionValidator newState,
-                             final ThreadStateTransitionValidator oldState) {
-            ++numChanges;
-            if (this.newState != null) {
-                if (this.newState != oldState) {
-                    throw new RuntimeException("State mismatch " + oldState + " different from " + this.newState);
-                }
-            }
-            this.oldState = oldState;
-            this.newState = newState;
-        }
+        internalTopologyBuilder.addSource(null, "name", null, null, null, topic1);
+        internalTopologyBuilder.addSink("out", "output", null, null, null, "name");
+
+        thread.setState(StreamThread.State.STARTING);
+        thread.rebalanceListener.onPartitionsRevoked(Collections.emptySet());
+
+        final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
+        final List<TopicPartition> assignedPartitions = new ArrayList<>();
+
+        // assign single partition
+        assignedPartitions.add(t1p1);
+        activeTasks.put(task1, Collections.singleton(t1p1));
+
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
+
+        final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
+        mockConsumer.assign(assignedPartitions);
+        mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
+        thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
+
+        thread.runOnce();
+
+        assertThat(thread.activeTasks().size(), equalTo(1));
+
+        thread.rebalanceListener.onPartitionsRevoked(assignedPartitions);
+        assertTrue(clientSupplier.producers.get(0).transactionCommitted());
+        assertFalse(clientSupplier.producers.get(0).closed());
+        assertEquals(1, thread.activeTasks().size());
     }
 
-    @Ignore
     @Test
     public void shouldReturnActiveTaskMetadataWhileRunningState() {
         internalTopologyBuilder.addSource(null, "source", null, null, null, topic1);
@@ -1043,16 +1038,12 @@ public void shouldReturnActiveTaskMetadataWhileRunningState() {
 
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
         final List<TopicPartition> assignedPartitions = new ArrayList<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
 
         // assign single partition
         assignedPartitions.add(t1p1);
         activeTasks.put(task1, Collections.singleton(t1p1));
-        partitionsToTaskId.put(t1p1, task1);
 
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         final MockConsumer<byte[], byte[]> mockConsumer = (MockConsumer<byte[], byte[]>) thread.consumer;
         mockConsumer.assign(assignedPartitions);
@@ -1076,7 +1067,6 @@ public void shouldReturnActiveTaskMetadataWhileRunningState() {
         assertEquals(CLIENT_ID + "-admin", metadata.adminClientId());
     }
 
-    @Ignore
     @Test
     public void shouldReturnStandbyTaskMetadataWhileRunningState() {
         internalStreamsBuilder.stream(Collections.singleton(topic1), consumed)
@@ -1087,7 +1077,7 @@ public void shouldReturnStandbyTaskMetadataWhileRunningState() {
         final MockConsumer<byte[], byte[]> restoreConsumer = clientSupplier.restoreConsumer;
         restoreConsumer.updatePartitions(
             "stream-thread-test-count-one-changelog",
-            singletonList(
+            Collections.singletonList(
                 new PartitionInfo("stream-thread-test-count-one-changelog",
                                   0,
                                   null,
@@ -1105,15 +1095,11 @@ public void shouldReturnStandbyTaskMetadataWhileRunningState() {
         thread.rebalanceListener.onPartitionsRevoked(Collections.emptySet());
 
         final Map<TaskId, Set<TopicPartition>> standbyTasks = new HashMap<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
 
         // assign single partition
         standbyTasks.put(task1, Collections.singleton(t1p1));
-        partitionsToTaskId.put(t1p1, task1);
 
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(Collections.emptyMap(), standbyTasks);
+        thread.taskManager().handleAssignment(Collections.emptyMap(), standbyTasks);
 
         thread.rebalanceListener.onPartitionsAssigned(Collections.emptyList());
 
@@ -1126,7 +1112,6 @@ public void shouldReturnStandbyTaskMetadataWhileRunningState() {
     }
 
     @SuppressWarnings("unchecked")
-    @Ignore
     @Test
     public void shouldUpdateStandbyTask() throws Exception {
         final String storeName1 = "count-one";
@@ -1147,69 +1132,62 @@ public void shouldUpdateStandbyTask() throws Exception {
         final StreamThread thread = createStreamThread(CLIENT_ID, config, false);
         final MockConsumer<byte[], byte[]> restoreConsumer = clientSupplier.restoreConsumer;
         restoreConsumer.updatePartitions(changelogName1,
-                                         singletonList(
-                                             new PartitionInfo(
-                                                 changelogName1,
-                                                 1,
-                                                 null,
-                                                 new Node[0],
-                                                 new Node[0]
-                                             )
-                                         )
+            Collections.singletonList(new PartitionInfo(changelogName1, 1, null, new Node[0], new Node[0]))
         );
 
-        restoreConsumer.assign(Utils.mkSet(partition1, partition2));
         restoreConsumer.updateEndOffsets(Collections.singletonMap(partition1, 10L));
         restoreConsumer.updateBeginningOffsets(Collections.singletonMap(partition1, 0L));
         restoreConsumer.updateEndOffsets(Collections.singletonMap(partition2, 10L));
         restoreConsumer.updateBeginningOffsets(Collections.singletonMap(partition2, 0L));
-        // let the store1 be restored from 0 to 10; store2 be restored from 5 (checkpointed) to 10
         final OffsetCheckpoint checkpoint
             = new OffsetCheckpoint(new File(stateDirectory.directoryForTask(task3), CHECKPOINT_FILE_NAME));
         checkpoint.write(Collections.singletonMap(partition2, 5L));
 
-        for (long i = 0L; i < 10L; i++) {
-            restoreConsumer.addRecord(new ConsumerRecord<>(
-                changelogName1,
-                1,
-                i,
-                ("K" + i).getBytes(),
-                ("V" + i).getBytes()));
-            restoreConsumer.addRecord(new ConsumerRecord<>(
-                changelogName2,
-                1,
-                i,
-                ("K" + i).getBytes(),
-                ("V" + i).getBytes()));
-        }
-
         thread.setState(StreamThread.State.STARTING);
         thread.rebalanceListener.onPartitionsRevoked(Collections.emptySet());
 
         final Map<TaskId, Set<TopicPartition>> standbyTasks = new HashMap<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
 
         // assign single partition
         standbyTasks.put(task1, Collections.singleton(t1p1));
         standbyTasks.put(task3, Collections.singleton(t2p1));
-        partitionsToTaskId.put(t1p1, task1);
 
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(Collections.emptyMap(), standbyTasks);
+        thread.taskManager().handleAssignment(Collections.emptyMap(), standbyTasks);
 
         thread.rebalanceListener.onPartitionsAssigned(Collections.emptyList());
 
         thread.runOnce();
 
-        final StandbyTask standbyTask1 = thread.taskManager().standbyTask(partition1);
-        final StandbyTask standbyTask2 = thread.taskManager().standbyTask(partition2);
+        final StandbyTask standbyTask1 = thread.taskManager().standbyTask(t1p1);
+        final StandbyTask standbyTask2 = thread.taskManager().standbyTask(t2p1);
+        assertEquals(task1, standbyTask1.id());
+        assertEquals(task3, standbyTask2.id());
+
         final KeyValueStore<Object, Long> store1 = (KeyValueStore<Object, Long>) standbyTask1.getStore(storeName1);
         final KeyValueStore<Object, Long> store2 = (KeyValueStore<Object, Long>) standbyTask2.getStore(storeName2);
+        assertEquals(0L, store1.approximateNumEntries());
+        assertEquals(0L, store2.approximateNumEntries());
+
+        // let the store1 be restored from 0 to 10; store2 be restored from 5 (checkpointed) to 10
+        for (long i = 0L; i < 10L; i++) {
+            restoreConsumer.addRecord(new ConsumerRecord<>(
+                changelogName1,
+                1,
+                i,
+                ("K" + i).getBytes(),
+                ("V" + i).getBytes()));
+            restoreConsumer.addRecord(new ConsumerRecord<>(
+                changelogName2,
+                1,
+                i,
+                ("K" + i).getBytes(),
+                ("V" + i).getBytes()));
+        }
+
+        thread.runOnce();
 
         assertEquals(10L, store1.approximateNumEntries());
-        assertEquals(5L, store2.approximateNumEntries());
-        assertEquals(0, thread.standbyRecords().size());
+        assertEquals(4L, store2.approximateNumEntries());
     }
 
     @Test
@@ -1243,33 +1221,6 @@ public void shouldNotCreateStandbyTaskIfStateStoresHaveLoggingDisabled() {
         assertThat(standbyTask, nullValue());
     }
 
-    private void setupInternalTopologyWithoutState() {
-        final MockProcessor mockProcessor = new MockProcessor();
-        internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
-        internalTopologyBuilder.addProcessor("processor1", () -> mockProcessor, "source1");
-    }
-
-    private StandbyTask createStandbyTask() {
-        final LogContext logContext = new LogContext("test");
-        final Logger log = logContext.logger(StreamThreadTest.class);
-        final StreamsMetricsImpl streamsMetrics =
-            new StreamsMetricsImpl(metrics, CLIENT_ID, StreamsConfig.METRICS_LATEST);
-        final StreamThread.StandbyTaskCreator standbyTaskCreator = new StreamThread.StandbyTaskCreator(
-            internalTopologyBuilder,
-            config,
-            streamsMetrics,
-            stateDirectory,
-            new MockChangelogReader(),
-            mockTime,
-            CLIENT_ID,
-            log);
-        return standbyTaskCreator.createTask(
-            new MockConsumer<>(OffsetResetStrategy.EARLIEST),
-            new TaskId(1, 2),
-            Collections.emptySet());
-    }
-
-    @Ignore
     @Test
     public void shouldPunctuateActiveTask() {
         final List<Long> punctuatedStreamTime = new ArrayList<>();
@@ -1299,16 +1250,12 @@ public void close() {}
         final List<TopicPartition> assignedPartitions = new ArrayList<>();
 
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
-        final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
 
         // assign single partition
         assignedPartitions.add(t1p1);
         activeTasks.put(task1, Collections.singleton(t1p1));
-        partitionsToTaskId.put(t1p1, task1);
 
-//        thread.taskManager().setPartitionsToTaskId(partitionsToTaskId);
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         clientSupplier.consumer.assign(assignedPartitions);
         clientSupplier.consumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L));
@@ -1362,7 +1309,6 @@ public void shouldAlwaysUpdateTasksMetadataAfterChangingState() {
         assertEquals(StreamThread.State.RUNNING.name(), metadata.threadState());
     }
 
-    @Ignore
     @Test
     public void shouldAlwaysReturnEmptyTasksMetadataWhileRebalancingStateAndTasksNotRunning() {
         internalStreamsBuilder.stream(Collections.singleton(topic1), consumed)
@@ -1406,16 +1352,22 @@ public void shouldAlwaysReturnEmptyTasksMetadataWhileRebalancingStateAndTasksNot
         activeTasks.put(task1, Collections.singleton(t1p1));
         standbyTasks.put(task2, Collections.singleton(t1p2));
 
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(activeTasks, standbyTasks);
+        thread.taskManager().handleAssignment(activeTasks, standbyTasks);
 
         thread.rebalanceListener.onPartitionsAssigned(assignedPartitions);
 
         assertThreadMetadataHasEmptyTasksWithState(thread.threadMetadata(), StreamThread.State.PARTITIONS_ASSIGNED);
     }
 
+    private void assertThreadMetadataHasEmptyTasksWithState(final ThreadMetadata metadata, final StreamThread.State state) {
+        assertEquals(state.name(), metadata.threadState());
+        assertTrue(metadata.activeTasks().isEmpty());
+        assertTrue(metadata.standbyTasks().isEmpty());
+    }
+
     @Ignore
     @Test
+    // FIXME: should unblock this test after we added invalid offset handling
     public void shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore() throws Exception {
         internalStreamsBuilder.stream(Collections.singleton("topic"), consumed)
                               .groupByKey().count(Materialized.as("count"));
@@ -1431,13 +1383,12 @@ public void shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore() t
         final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
         final TaskId task0 = new TaskId(0, 0);
         activeTasks.put(task0, topicPartitionSet);
-//        thread.taskManager().setPartitionsToTaskId(Collections.singletonMap(topicPartition, task0));
-        //FIXME
-//        thread.taskManager().setAssignmentMetadata(activeTasks, Collections.emptyMap());
+
+        thread.taskManager().handleAssignment(activeTasks, Collections.emptyMap());
 
         mockConsumer.updatePartitions(
             "topic",
-            singletonList(
+            Collections.singletonList(
                 new PartitionInfo(
                     "topic",
                     0,
@@ -1451,7 +1402,7 @@ public void shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore() t
 
         mockRestoreConsumer.updatePartitions(
             "stream-thread-test-count-changelog",
-            singletonList(
+            Collections.singletonList(
                 new PartitionInfo(
                     "stream-thread-test-count-changelog",
                     0,
@@ -1738,16 +1689,8 @@ private void verifyLogMessagesForSkippedRecordsForInvalidTimestamps(final LogCap
         ));
     }
 
-    private void assertThreadMetadataHasEmptyTasksWithState(final ThreadMetadata metadata,
-                                                            final StreamThread.State state) {
-        assertEquals(state.name(), metadata.threadState());
-        assertTrue(metadata.activeTasks().isEmpty());
-        assertTrue(metadata.standbyTasks().isEmpty());
-    }
-
     @Test
-    // TODO: Need to add a test case covering EOS when we create a mock taskManager class
-    public void producerMetricsVerificationWithoutEOS() {
+    public void shouldConstructProducerMetricsWithoutEOS() {
         final MockProducer<byte[], byte[]> producer = new MockProducer<>();
         final Consumer<byte[], byte[]> consumer = EasyMock.createNiceMock(Consumer.class);
         final TaskManager taskManager = mockTaskManagerCommit(consumer, 1, 0);
@@ -1783,7 +1726,45 @@ public void producerMetricsVerificationWithoutEOS() {
     }
 
     @Test
-    public void adminClientMetricsVerification() {
+    public void shouldConstructProducerMetricsWithEOS() {
+        final MockProducer<byte[], byte[]> producer = new MockProducer<>();
+        final Consumer<byte[], byte[]> consumer = EasyMock.createNiceMock(Consumer.class);
+        final TaskManager taskManager = mockTaskManagerCommit(consumer, 1, 0);
+
+        final StreamsMetricsImpl streamsMetrics =
+            new StreamsMetricsImpl(metrics, CLIENT_ID, StreamsConfig.METRICS_LATEST);
+        final StreamThread thread = new StreamThread(
+            mockTime,
+            new StreamsConfig(configProps(true)),
+            null,       // with EOS the thread producer should be null
+            null,
+            consumer,
+            consumer,
+            null,
+            null,
+            taskManager,
+            streamsMetrics,
+            internalTopologyBuilder,
+            CLIENT_ID,
+            new LogContext(""),
+            new AtomicInteger()
+        );
+        final MetricName testMetricName = new MetricName("test_metric", "", "", new HashMap<>());
+        final Metric testMetric = new KafkaMetric(
+            new Object(),
+            testMetricName,
+            (Measurable) (config, now) -> 0,
+            null,
+            new MockTime());
+
+        // without creating tasks the metrics should be empty
+        producer.setMockMetrics(testMetricName, testMetric);
+        final Map<MetricName, Metric> producerMetrics = thread.producerMetrics();
+        assertEquals(Collections.emptyMap(), producerMetrics);
+    }
+
+    @Test
+    public void shouldConstructAdminMetrics() {
         final Node broker1 = new Node(0, "dummyHost-1", 1234);
         final Node broker2 = new Node(1, "dummyHost-2", 1234);
         final List<Node> cluster = Arrays.asList(broker1, broker2);
@@ -1826,6 +1807,42 @@ public void adminClientMetricsVerification() {
         assertEquals(testMetricName, adminClientMetrics.get(testMetricName).metricName());
     }
 
+    private TaskManager mockTaskManagerCommit(final Consumer<byte[], byte[]> consumer,
+                                              final int numberOfCommits,
+                                              final int commits) {
+        final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);
+        EasyMock.expect(taskManager.fixmeStreamTasks()).andReturn(Collections.emptyMap()).anyTimes();
+        EasyMock.expect(taskManager.commitAll()).andReturn(commits).times(numberOfCommits);
+        EasyMock.replay(taskManager, consumer);
+        return taskManager;
+    }
+
+    private void setupInternalTopologyWithoutState() {
+        final MockProcessor mockProcessor = new MockProcessor();
+        internalTopologyBuilder.addSource(null, "source1", null, null, null, topic1);
+        internalTopologyBuilder.addProcessor("processor1", () -> mockProcessor, "source1");
+    }
+
+    private StandbyTask createStandbyTask() {
+        final LogContext logContext = new LogContext("test");
+        final Logger log = logContext.logger(StreamThreadTest.class);
+        final StreamsMetricsImpl streamsMetrics =
+            new StreamsMetricsImpl(metrics, CLIENT_ID, StreamsConfig.METRICS_LATEST);
+        final StreamThread.StandbyTaskCreator standbyTaskCreator = new StreamThread.StandbyTaskCreator(
+            internalTopologyBuilder,
+            config,
+            streamsMetrics,
+            stateDirectory,
+            new MockChangelogReader(),
+            mockTime,
+            CLIENT_ID,
+            log);
+        return standbyTaskCreator.createTask(
+            new MockConsumer<>(OffsetResetStrategy.EARLIEST),
+            new TaskId(1, 2),
+            Collections.emptySet());
+    }
+
     private void addRecord(final MockConsumer<byte[], byte[]> mockConsumer,
                            final long offset) {
         addRecord(mockConsumer, offset, -1L);
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
index e57c8f78b68d8..5e48f21c3d322 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
@@ -25,9 +25,9 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.internals.KafkaFutureImpl;
-import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.processor.TaskId;
+
 import org.easymock.EasyMock;
 import org.easymock.EasyMockRunner;
 import org.easymock.Mock;
@@ -267,7 +267,7 @@ public void shouldSuspendActiveTasks() {
     }
 
     @Test
-    public void shouldThrowStreamsExceptionAtEndIfExceptionDuringSuspend() {
+    public void shouldPassUpIfExceptionDuringSuspend() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true) {
             @Override
             public void suspend() {
@@ -283,7 +283,7 @@ public void suspend() {
         assertThat(taskManager.checkForCompletedRestoration(), is(true));
         assertThat(task00.state(), is(Task.State.RUNNING));
 
-        assertThrows(StreamsException.class, () -> taskManager.handleRevocation(taskId00Partitions));
+        assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(taskId00Partitions));
         assertThat(task00.state(), is(Task.State.RUNNING));
     }
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/StreamThreadStateStoreProviderTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/StreamThreadStateStoreProviderTest.java
index 452f89a5e1669..fc1ae7a0e543a 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/StreamThreadStateStoreProviderTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/StreamThreadStateStoreProviderTest.java
@@ -64,6 +64,7 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.instanceOf;
@@ -343,7 +344,7 @@ private StreamTask createStreamsTask(final StreamsConfig streamsConfig,
     private void mockThread(final boolean initialized) {
         EasyMock.expect(threadMock.isRunning()).andReturn(initialized);
         EasyMock.expect(threadMock.allTasks()).andStubReturn(tasks);
-        EasyMock.expect(threadMock.activeTasks()).andStubReturn(tasks.values().stream()::iterator);
+        EasyMock.expect(threadMock.activeTasks()).andStubReturn(tasks.values().stream().collect(Collectors.toList()));
         EasyMock.expect(threadMock.state()).andReturn(
             initialized ? StreamThread.State.RUNNING : StreamThread.State.PARTITIONS_ASSIGNED
         ).anyTimes();
diff --git a/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java b/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java
index 159ca4644c832..54e868eb4c2ea 100644
--- a/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/tests/StreamsUpgradeTest.java
@@ -23,7 +23,6 @@
 import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.RebalanceProtocol;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.common.Cluster;
-import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.utils.ByteBufferInputStream;
@@ -43,7 +42,6 @@
 import org.apache.kafka.streams.processor.internals.assignment.AssignorError;
 import org.apache.kafka.streams.processor.internals.assignment.LegacySubscriptionInfoSerde;
 import org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo;
-import org.apache.kafka.streams.state.HostInfo;
 
 import java.io.ByteArrayOutputStream;
 import java.io.DataInputStream;
@@ -91,7 +89,7 @@ public static void main(final String[] args) throws Exception {
         });
     }
 
-    public static KafkaStreams buildStreams(final Properties streamsProperties) {
+    private static KafkaStreams buildStreams(final Properties streamsProperties) {
         final StreamsBuilder builder = new StreamsBuilder();
         final KStream<Void, Void> dataStream = builder.stream("data");
         dataStream.process(SmokeTestUtil.printProcessorSupplier("data"));
@@ -213,26 +211,10 @@ public void onAssignment(final ConsumerPartitionAssignor.Assignment assignment,
             final List<TopicPartition> partitions = new ArrayList<>(assignment.partitions());
             partitions.sort(PARTITION_COMPARATOR);
 
-            // version 1 field
-            final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
-            // version 2 fields
-            final Map<TopicPartition, PartitionInfo> topicToPartitionInfo = new HashMap<>();
-            final Map<HostInfo, Set<TopicPartition>> partitionsByHost;
-
-            final Map<TopicPartition, TaskId> partitionsToTaskId = new HashMap<>();
-
-            processVersionTwoAssignment("test ", info, partitions, activeTasks, topicToPartitionInfo, partitionsToTaskId);
-            partitionsByHost = info.partitionsByHost();
+            final Map<TaskId, Set<TopicPartition>> activeTasks = getActiveTasks(partitions, info);
 
             final TaskManager taskManager = taskManger();
-// TODO K9113: fix this test
-//            taskManager.setClusterMetadata(Cluster.empty().withPartitions(topicToPartitionInfo));
-//            taskManager.setPartitionsByHostState(partitionsByHost);
-//            taskManager.setPartitionsToTaskId(partitionsToTaskId);
-//            taskManager.setAssignmentMetadata(activeTasks, info.standbyTasks());
-            taskManager.fixmeUpdateSubscriptionsFromAssignment(partitions);
-            // TODO K9113: fix this test
-            //  taskManager.handleRebalanceStart(false);
+            taskManager.handleAssignment(activeTasks, info.standbyTasks());
             usedSubscriptionMetadataVersionPeek.set(usedSubscriptionMetadataVersion);
         }
 
@@ -265,8 +247,8 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr
                     final Subscription subscription = entry.getValue();
 
                     final SubscriptionInfo info = SubscriptionInfo.decode(subscription.userData()
-                                                                                      .putInt(0, LATEST_SUPPORTED_VERSION)
-                                                                                      .putInt(4, LATEST_SUPPORTED_VERSION));
+                        .putInt(0, LATEST_SUPPORTED_VERSION)
+                        .putInt(4, LATEST_SUPPORTED_VERSION));
 
                     downgradedSubscriptions.put(
                         entry.getKey(),
@@ -329,7 +311,7 @@ private static class FutureSubscriptionInfo {
             }
         }
 
-        public ByteBuffer encode() {
+        private ByteBuffer encode() {
             final ByteBuffer buf = encodeFutureVersion();
             buf.rewind();
             return buf;
diff --git a/streams/src/test/java/org/apache/kafka/test/MockRecordCollector.java b/streams/src/test/java/org/apache/kafka/test/MockRecordCollector.java
index 597a8b771d9c9..a0fdb72d4f206 100644
--- a/streams/src/test/java/org/apache/kafka/test/MockRecordCollector.java
+++ b/streams/src/test/java/org/apache/kafka/test/MockRecordCollector.java
@@ -42,9 +42,6 @@ public class MockRecordCollector implements RecordCollector {
     // remember if flushed is called
     private boolean flushed = false;
 
-    @Override
-    public void initialize() {}
-
     @Override
     public <K, V> void send(final String topic,
                             final K key,