如何为批量刷新的递归逻辑编写正确的单元测试

场景与问题概述 #

我有一个循环执行的逻辑：通过refreshInBatches()查询数据库返回nodeId，用Set记录已访问节点，终止条件为到达id=1的节点或遇到已访问过的节点。现在需要写单元测试验证该逻辑，预期refreshInBatches()会被调用5次（模拟的响应列表有5个元素，最后一个是结束节点），但实际仅调用2次，因为每次Mock返回的nodeId都是第一个值10L，导致循环提前终止。

业务逻辑代码 #

public Mono<MyClass> refresh(String databaseName, Long nodeId, Integer batchSize, Set<Long> visitedNodes) {
    Mono<MyClass> result = null;
    do {
      addNodeToSet(visitedNodes, result);
      result = refreshInBatches(databaseName, nodeId, batchSize);
    } while (notEndNode(nodeId, result.block(), visitedNodes));
    return result;
}

测试代码（原错误版本） #

@InjectMocks
DbRepository dbRepository = Mockito.spy(new DbRepository(dbClient, null));

MyClass response = Mockito.spy(new MyClass());

@Mock
Node node = Mockito.mock(Node.class);

@Test
public void testRefreshDBShouldCallTheRefreshMethodFiveTimes() {

    List<Value> setIds = Arrays.asList(Values.value(10L), Values.value(20L), Values.value(50L), Values.value(60L), Values.value(1L));
    Iterator<Value> iterator = setIds.iterator();
    Mockito.doReturn(iterator.next()).when(node).get(anyString());
    Mockito.doReturn(node).when(response).getNode();
    Mockito.doReturn(Mono.just(response)).when(dbRepository).refreshInBatches(anyString(), anyLong(), anyInt());
    dbRepository.refresh("dataBase", -1L, 10000, new HashSet<>());
    verify(dbRepository, times(setIds.size())).refreshInBatches(anyString(), anyLong(), anyInt());
}

辅助方法 #

private void addNodeToSet(Set<Long> visited, Mono<MyClass> projection) {
    if (ObjectUtils.isNotEmpty(projection) && ObjectUtils.isNotEmpty(projection.block())) {
      visited.add(projection.block().getNode().get("id").asLong());
    }
}

private static boolean notEndNode(Long nodeId, MyClass projection, Set<Long> visited) {
    return projection != null &&
        ObjectUtils.isNotEmpty(projection.getNode()) &&
        !nodeId.equals(projection.getNode().get("id").asLong()) &&
        !visited.contains(projection.getNode().get("id").asLong());
}

问题原因 #

1. Mockito静态绑定值：Mockito.doReturn(iterator.next()).when(node).get(anyString())是静态绑定逻辑，初始化时就执行iterator.next()拿到第一个值10L，之后每次调用node.get()都会返回这个固定值，不会再触发iterator.next()获取下一个元素。
2. 复用同一个响应实例：测试中始终返回同一个response对象，其内部的node id不会变化，导致notEndNode判断很快不满足，循环提前终止。

解决方案 #

1. 用thenAnswer动态生成Mock返回值 #

thenAnswer会在每次Mock方法被调用时执行自定义逻辑，能实现每次调用node.get()都获取迭代器的下一个值。

2. 每次调用生成新的响应实例 #

确保每次refreshInBatches返回的MyClass实例都绑定当前的node id，避免复用同一个实例导致值不变。

修改后的测试代码：

@InjectMocks
DbRepository dbRepository = Mockito.spy(new DbRepository(dbClient, null));

@Mock
Node node;

@Test
public void testRefreshDBShouldCallTheRefreshMethodFiveTimes() {
    List<Value> setIds = Arrays.asList(Values.value(10L), Values.value(20L), Values.value(50L), Values.value(60L), Values.value(1L));
    Iterator<Value> iterator = setIds.iterator();

    // 动态返回迭代器的下一个值
    Mockito.when(node.get(anyString())).thenAnswer(invocation -> {
        if (iterator.hasNext()) {
            return iterator.next();
        }
        throw new NoSuchElementException();
    });

    // 每次调用refreshInBatches都生成新的MyClass实例
    Mockito.when(dbRepository.refreshInBatches(anyString(), anyLong(), anyInt()))
            .thenAnswer(invocation -> {
                MyClass response = new MyClass();
                Mockito.when(response.getNode()).thenReturn(node);
                return Mono.just(response);
            });

    dbRepository.refresh("dataBase", -1L, 10000, new HashSet<>());
    verify(dbRepository, times(setIds.size())).refreshInBatches(anyString(), anyLong(), anyInt());
}

额外优化：业务代码避免重复block #

原业务代码中多次调用result.block()会导致Mono被重复订阅，建议修改为只block一次并复用结果：

public Mono<MyClass> refresh(String databaseName, Long nodeId, Integer batchSize, Set<Long> visitedNodes) {
    Mono<MyClass> result = null;
    MyClass currentResult = null;
    do {
        addNodeToSet(visitedNodes, currentResult);
        result = refreshInBatches(databaseName, nodeId, batchSize);
        currentResult = result.block();
    } while (notEndNode(nodeId, currentResult, visitedNodes));
    return result;
}

// 对应修改addNodeToSet，直接传MyClass实例避免重复block
private void addNodeToSet(Set<Long> visited, MyClass projection) {
    if (ObjectUtils.isNotEmpty(projection)) {
        visited.add(projection.getNode().get("id").asLong());
    }
}

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