如何对使用leaderelection.RunOrDie的K8s选主函数进行单元测试

首先先明确你的核心需求：通过client-go/fake客户端，确定性地控制Leader选举的状态，验证SimpleLeader的amLeader值是否符合预期。以下是具体的实现思路和示例代码：

测试核心逻辑 #

K8s Leader选举的本质是基于Lease资源的持有者身份判断，因此我们可以通过fake client预先注入或动态修改Lease对象的状态，来模拟不同的选举场景，进而验证代码的回调逻辑是否正确触发。

具体实现步骤 #

1. 初始化测试依赖 #

• 使用fake.NewSimpleClientset()创建测试用的fake客户端
• 初始化SimpleLeader实例
• 创建可取消的context，用于控制run方法的启停

2. 模拟「成功当选主节点」场景 #

• 预先通过fake client创建Lease资源，将其spec.holderIdentity设置为当前实例的标识testidentity
• 启动run方法的goroutine
• 等待回调触发后，断言amLeader.Load()返回true

3. 模拟「未当选主节点」场景 #

• 预先创建Lease资源，将spec.holderIdentity设置为其他节点标识（比如other-node）
• 启动run方法goroutine
• 断言amLeader.Load()保持false

4. 模拟「丢失主节点身份」场景 #

• 先让实例成功当选（按步骤2操作）
• 在run方法运行时，通过fake client修改Lease的spec.holderIdentity为其他值
• 等待OnStoppedLeading回调触发，断言amLeader.Load()变为false

示例测试代码 #

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes/fake"
	coordinationv1 "k8s.io/api/coordination/v1"
)

// 辅助函数：生成指针类型，适配LeaseSpec的字段要求
func pointer[T any](v T) *T {
	return &v
}

func TestSimpleLeader_ElectionScenarios(t *testing.T) {
	// 测试场景1：成功当选主节点
	t.Run("successfully become leader", func(t *testing.T) {
		fakeClient := fake.NewSimpleClientset()
		sl := &SimpleLeader{}
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		// 预先创建当前实例持有的Lease
		lease := &coordinationv1.Lease{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "testname",
				Namespace: "testnamespace",
			},
			Spec: coordinationv1.LeaseSpec{
				HolderIdentity:       pointer.String("testidentity"),
				LeaseDurationSeconds: pointer.Int32(1), // 缩短超时加快测试
			},
		}
		_, err := fakeClient.CoordinationV1().Leases("testnamespace").Create(ctx, lease, metav1.CreateOptions{})
		assert.NoError(t, err)

		// 启动Leader选举逻辑
		go sl.run(ctx, fakeClient.CoordinationV1())

		// 等待回调触发（生产级测试可改用sync.Cond替代Sleep）
		time.Sleep(100 * time.Millisecond)
		isLeader, ok := sl.amLeader.Load().(bool)
		assert.True(t, ok)
		assert.True(t, isLeader)
	})

	// 测试场景2：无法当选主节点（Lease被其他实例持有）
	t.Run("fail to become leader when lease is held by others", func(t *testing.T) {
		fakeClient := fake.NewSimpleClientset()
		sl := &SimpleLeader{}
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		// 创建被其他实例持有的Lease
		lease := &coordinationv1.Lease{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "testname",
				Namespace: "testnamespace",
			},
			Spec: coordinationv1.LeaseSpec{
				HolderIdentity:       pointer.String("other-node"),
				LeaseDurationSeconds: pointer.Int32(1),
			},
		}
		_, err := fakeClient.CoordinationV1().Leases("testnamespace").Create(ctx, lease, metav1.CreateOptions{})
		assert.NoError(t, err)

		go sl.run(ctx, fakeClient.CoordinationV1())

		time.Sleep(300 * time.Millisecond)
		isLeader, ok := sl.amLeader.Load().(bool)
		assert.True(t, ok)
		assert.False(t, isLeader)
	})

	// 测试场景3：丢失主节点身份
	t.Run("lose leader status when lease is stolen", func(t *testing.T) {
		fakeClient := fake.NewSimpleClientset()
		sl := &SimpleLeader{}
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		// 先让实例当选主节点
		lease := &coordinationv1.Lease{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "testname",
				Namespace: "testnamespace",
			},
			Spec: coordinationv1.LeaseSpec{
				HolderIdentity:       pointer.String("testidentity"),
				LeaseDurationSeconds: pointer.Int32(1),
			},
		}
		_, err := fakeClient.CoordinationV1().Leases("testnamespace").Create(ctx, lease, metav1.CreateOptions{})
		assert.NoError(t, err)

		go sl.run(ctx, fakeClient.CoordinationV1())
		time.Sleep(100 * time.Millisecond)
		isLeader, ok := sl.amLeader.Load().(bool)
		assert.True(t, ok)
		assert.True(t, isLeader)

		// 修改Lease持有者为其他节点
		lease.Spec.HolderIdentity = pointer.String("other-node")
		_, err = fakeClient.CoordinationV1().Leases("testnamespace").Update(ctx, lease, metav1.UpdateOptions{})
		assert.NoError(t, err)

		// 等待回调触发，验证状态切换
		time.Sleep(150 * time.Millisecond)
		isLeader, ok = sl.amLeader.Load().(bool)
		assert.True(t, ok)
		assert.False(t, isLeader)
	})
}

额外注意事项 #

• 示例中用time.Sleep是为了简化演示，生产级测试推荐使用sync.Cond或通道来等待回调触发，避免依赖不确定的时间
• 可以修改LeaderElectionConfig中的超时参数（比如缩短LeaseDuration和RenewDeadline），加快测试执行速度
• atomic.Value读取时需要断言类型，确保类型安全

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