Golang指数退避重试函数的单元测试实现方法

要测试这个指数退避重试的等待时间，核心思路是mock依赖函数，避免真实的sleep等待，同时记录每次等待的时长来验证是否符合预期。下面分步骤说明实现方式：

1. 先重构原函数，使其可测试 #

原函数直接硬编码调用了ssh.NewSignerFromKey和time.Sleep，还依赖包级变量privateKeyMaterial，这些都不利于测试。我们需要把依赖抽象出来，方便测试时替换：

import (
    "math"
    "time"
    "golang.org/x/crypto/ssh"
)

// 定义依赖结构体，封装外部调用
type odinDeps struct {
    newSignerFunc func(key any) (ssh.Signer, error)
    sleepFunc     func(d time.Duration)
}

// 包级依赖变量，测试时可替换
var odinDeps = odinDeps{
    newSignerFunc: ssh.NewSignerFromKey,
    sleepFunc:     time.Sleep,
}

func (oc OdinClient) GetPrivateKeySigner(materialSet string) (ssh.Signer, error) {
    signer, err := odinDeps.newSignerFunc(privateKeyMaterial.PrivateKey)
    if err != nil {
        for i := 0; i < 5; i++ {
            waitMs := time.Duration(math.Pow(2, float64(i)) * 100) * time.Millisecond
            odinDeps.sleepFunc(waitMs)
            
            signer, err := odinDeps.newSignerFunc(privateKeyMaterial.PrivateKey)
            if err == nil {
                return signer, nil
            }
        }
        return nil, novaerr.New("utils/getPrivateKey", "failed", err)
    }
    return signer, nil
}

2. 编写单元测试 #

我们需要两个核心测试场景：

• 场景1：所有重试都失败，验证完整的指数退避等待序列
• 场景2：某次重试成功，验证只执行到对应次数的等待

测试代码示例 #

package yourpackage

import (
    "testing"
    "time"
    "github.com/stretchr/testify/assert"
    "golang.org/x/crypto/ssh"
)

// 测试完整的指数退避等待序列（5次重试全部失败）
func TestGetPrivateKeySigner_FullBackoff(t *testing.T) {
    // 保存原始依赖，测试后恢复，避免影响其他用例
    originalDeps := odinDeps
    defer func() { odinDeps = originalDeps }()

    // 记录每次sleep的时长
    var sleepRecords []time.Duration
    // mock sleep函数，只记录时长不真实等待
    odinDeps.sleepFunc = func(d time.Duration) {
        sleepRecords = append(sleepRecords, d)
    }

    // mock签名函数，每次都返回错误，触发全部重试
    odinDeps.newSignerFunc = func(key any) (ssh.Signer, error) {
        return nil, assert.AnError
    }

    // 调用目标函数
    oc := OdinClient{}
    _, err := oc.GetPrivateKeySigner("test-material")
    assert.Error(t, err)

    // 验证等待次数：5次重试对应5次sleep
    assert.Len(t, sleepRecords, 5)

    // 验证每次等待时长符合指数退避规则：100ms * 2^i（i从0到4）
    expectedDurations := []time.Duration{
        100 * time.Millisecond,
        200 * time.Millisecond,
        400 * time.Millisecond,
        800 * time.Millisecond,
        1600 * time.Millisecond,
    }
    for idx, actual := range sleepRecords {
        assert.Equal(t, expectedDurations[idx], actual)
    }
}

// 测试中途重试成功的场景（第3次尝试成功）
func TestGetPrivateKeySigner_SucceedsOnThirdAttempt(t *testing.T) {
    originalDeps := odinDeps
    defer func() { odinDeps = originalDeps }()

    var sleepRecords []time.Duration
    odinDeps.sleepFunc = func(d time.Duration) {
        sleepRecords = append(sleepRecords, d)
    }

    // 模拟前2次失败（首次尝试+第一次重试），第3次尝试成功
    attemptCount := 0
    odinDeps.newSignerFunc = func(key any) (ssh.Signer, error) {
        attemptCount++
        if attemptCount <= 2 {
            return nil, assert.AnError
        }
        // 返回mock的Signer实现
        return &mockSigner{}, nil
    }

    oc := OdinClient{}
    signer, err := oc.GetPrivateKeySigner("test-material")
    assert.NoError(t, err)
    assert.NotNil(t, signer)

    // 验证sleep次数：2次（对应前两次失败后的等待）
    assert.Len(t, sleepRecords, 2)
    assert.Equal(t, 100*time.Millisecond, sleepRecords[0])
    assert.Equal(t, 200*time.Millisecond, sleepRecords[1])
    // 验证总尝试次数：3次（首次+2次重试）
    assert.Equal(t, 3, attemptCount)
}

// Mock ssh.Signer的实现，满足接口要求
type mockSigner struct{}

func (m *mockSigner) PublicKey() ssh.PublicKey {
    return nil
}

func (m *mockSigner) Sign(_ ssh.Algorithm, rand io.Reader, data []byte) (*ssh.Signature, error) {
    return nil, nil
}

关键说明 #

• 依赖替换：通过把ssh.NewSignerFromKey和time.Sleep抽象为可替换的函数变量，我们可以在测试中完全控制这些依赖的行为，不需要真实等待或生成有效的SSH密钥。
• 验证逻辑：通过记录每次sleep的时长，直接对比是否符合100ms * 2^i的指数退避规则，确保等待时间正确。
• 场景覆盖：覆盖了全部重试失败和中途成功两种场景，保证函数在不同情况下的退避逻辑都符合预期。

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