@Test
public void testMaybeClean() throws Exception {

    File baseDir = Files.createTempDirectory("test").toFile();
    try {
        final long cleanupDelay = 1000L;
        Properties props = configProps();
        props.setProperty(StreamsConfig.STATE_CLEANUP_DELAY_MS_CONFIG, Long.toString(cleanupDelay));
        props.setProperty(StreamsConfig.STATE_DIR_CONFIG, baseDir.getCanonicalPath());
        StreamsConfig config = new StreamsConfig(props);
        File applicationDir = new File(baseDir, applicationId);
        applicationDir.mkdir();
        File stateDir1 = new File(applicationDir, task1.toString());
        File stateDir2 = new File(applicationDir, task2.toString());
        File stateDir3 = new File(applicationDir, task3.toString());
        File extraDir = new File(applicationDir, "X");
        stateDir1.mkdir();
        stateDir2.mkdir();
        stateDir3.mkdir();
        extraDir.mkdir();
        MockTime mockTime = new MockTime();
        TopologyBuilder builder = new TopologyBuilder();
        builder.addSource("source1", "topic1");
        StreamThread thread = new StreamThread(builder, config, new MockClientSupplier(), applicationId, clientId,  processId, new Metrics(), mockTime) {
            @Override
            public void maybeClean() {
                super.maybeClean();
            }
            @Override
            protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) {
                ProcessorTopology topology = builder.build("X", id.topicGroupId);
                return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config);
            }
        };
        initPartitionGrouper(config, thread);
        ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
        assertTrue(thread.tasks().isEmpty());
        mockTime.sleep(cleanupDelay);
        // all directories exist since an assignment didn't happen
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertTrue(stateDir3.exists());
        assertTrue(extraDir.exists());
        List<TopicPartition> revokedPartitions;
        List<TopicPartition> assignedPartitions;
        Map<Integer, StreamTask> prevTasks;
        //
        // Assign t1p1 and t1p2. This should create task1 & task2
        //
        revokedPartitions = Collections.emptyList();
        assignedPartitions = Arrays.asList(t1p1, t1p2);
        prevTasks = new HashMap(thread.tasks());
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        // there shouldn't be any previous task
        assertTrue(prevTasks.isEmpty());
        // task 1 & 2 are created
        assertEquals(2, thread.tasks().size());
        // all directories should still exit before the cleanup delay time
        mockTime.sleep(cleanupDelay - 10L);
        thread.maybeClean();
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertTrue(stateDir3.exists());
        assertTrue(extraDir.exists());
        // all state directories except for task task2 & task3 will be removed. the extra directory should still exists
        mockTime.sleep(11L);
        thread.maybeClean();
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
        //
        // Revoke t1p1 and t1p2. This should remove task1 & task2
        //
        revokedPartitions = assignedPartitions;
        assignedPartitions = Collections.emptyList();
        prevTasks = new HashMap(thread.tasks());
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        // previous tasks should be committed
        assertEquals(2, prevTasks.size());
        for (StreamTask task : prevTasks.values()) {
            assertTrue(((TestStreamTask) task).committed);
            ((TestStreamTask) task).committed = false;
        }
        // no task
        assertTrue(thread.tasks().isEmpty());
        // all state directories for task task1 & task2 still exist before the cleanup delay time
        mockTime.sleep(cleanupDelay - 10L);
        thread.maybeClean();
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
        // all state directories for task task1 & task2 are removed
        mockTime.sleep(11L);
        thread.maybeClean();
        assertFalse(stateDir1.exists());
        assertFalse(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
    } finally {
        Utils.delete(baseDir);
    }
}