如何让单元测试等待线程池任务执行完成后再退出

你可以用这几种方法让单元测试等线程池任务全做完再退出：

• 方法一：用shutdown() + awaitTermination()组合
  先调用executor.shutdown()关闭线程池，拒绝接收新任务；再调用awaitTermination()让单元测试的主线程阻塞，直到所有已提交的任务执行完毕，或者达到设置的超时时间。

修改后的测试方法代码：

@Test
public void test() throws InterruptedException {
    ExecutorService executor = Executors.newFixedThreadPool(5);

    for (int i = 0; i < 10; i++) {
        executor.submit(new MyTask(i));
    }

    executor.shutdown(); // 关闭线程池，不再接受新任务
    // 设置超时时间要长于任务的最大执行时长，这里设1小时10分钟确保覆盖任务的1小时休眠
    if (!executor.awaitTermination(1, java.util.concurrent.TimeUnit.HOURS)) {
        // 超时后强制关闭线程池，中断未完成的任务
        executor.shutdownNow();
    }
}

• 方法二：收集Future对象逐个等待
  submit()方法会返回Future对象，我们把所有任务的Future存入列表，然后遍历调用get()方法，每个get()会阻塞到对应任务执行完成。

修改后的测试方法代码：

@Test
public void test() throws InterruptedException, java.util.concurrent.ExecutionException {
    ExecutorService executor = Executors.newFixedThreadPool(5);
    java.util.List<java.util.concurrent.Future<?>> futures = new java.util.ArrayList<>();

    for (int i = 0; i < 10; i++) {
        futures.add(executor.submit(new MyTask(i)));
    }

    // 逐个等待任务完成
    for (java.util.concurrent.Future<?> future : futures) {
        future.get(); // 阻塞直到当前任务执行完毕
    }

    executor.shutdown(); // 所有任务完成后关闭线程池
}

这种方式还能处理单个任务抛出的异常（get()会抛出ExecutionException），方便排查问题。

• 方法三：使用CountDownLatch
  初始化一个计数器，数值等于任务总数；每个任务执行完毕后调用countDown()让计数器减1，主线程调用await()阻塞直到计数器归0。

修改后的完整代码：

package com.dolphin.client;

import lombok.extern.slf4j.Slf4j;
import org.junit.Test;

import java.time.Duration;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@Slf4j
public class FeignClientApplication {
    @Test
    public void test() throws InterruptedException {
        int taskCount = 10;
        CountDownLatch latch = new CountDownLatch(taskCount);
        ExecutorService executor = Executors.newFixedThreadPool(5);

        for (int i = 0; i < taskCount; i++) {
            executor.submit(new MyTask(i, latch));
        }

        latch.await(); // 等待所有任务完成
        executor.shutdown();
    }

    static class MyTask implements Runnable {
        private final int taskId;
        private final CountDownLatch latch;

        public MyTask(int taskId, CountDownLatch latch) {
            this.taskId = taskId;
            this.latch = latch;
        }

        @Override
        public void run() {
            try {
                Thread.sleep(Duration.ofHours(1));
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            System.out.println("Task ID : " + taskId + " performed by " + Thread.currentThread().getName());
            latch.countDown(); // 任务完成，计数器减1
        }
    }
}

这种方式适合需要精确控制等待任务数量的场景，灵活性较高。

内容的提问来源于stack exchange，提问作者Dolphin