Repository: pingcap/tidb
File: br/pkg/utils/register_test.go (Go)

Code before:
	testEtcdCluster := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1, GRPCKeepAliveInterval: time.Second, GRPCKeepAliveTimeout: 10 * time.Second})
	defer testEtcdCluster.Terminate(t)

	ctx := context.Background()
	client := testEtcdCluster.RandClient()

	register := NewTaskRegisterWithTTL(client, 3*time.Second, RegisterRestore, "test")
	defer func() {
		err := register.Close(ctx)
		// for flaky test, the lease would expire
		if err != nil && !strings.Contains(err.Error(), "requested lease not found") {
			require.NoError(t, err)
		}
	}()
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval", "return(200)"))
	defer func() {
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval")
	}()
	err := register.RegisterTask(ctx)
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		list, err := GetImportTasksFrom(ctx, client)
		if err != nil {
			// The task lease might be revoked during TTL query.
			return strings.Contains(err.Error(), "requested lease not found")
		}
		return len(list.Tasks) == 0
	}, 5*time.Second, 50*time.Millisecond)

	failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop")
	failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant")
	require.Eventually(t, func() bool {
		list, err := GetImportTasksFrom(ctx, client)
		return err == nil && len(list.Tasks) > 0
	}, 5*time.Second, 50*time.Millisecond)
	list, err := GetImportTasksFrom(ctx, client)
	require.NoError(t, err)
	for _, task := range list.Tasks {
		t.Log(task.MessageToUser())
		require.Equal(t, "/tidb/brie/import/restore/test", task.Key)
	}
	require.True(t, len(list.Tasks) > 0)
}


Review comment (security):
The assertions treat "requested lease not found" by matching a substring in `err.Error()`. This is sensitive to upstream etcd error message changes and can make the test flaky again. Prefer checking the concrete etcd error (e.g. `rpctypes.ErrLeaseNotFound`, possibly after unwrapping) rather than string matching.

Code after:
	defer testEtcdCluster.Terminate(t)

	ctx := context.Background()
	client := testEtcdCluster.RandClient()

	register := NewTaskRegisterWithTTL(client, 3*time.Second, RegisterRestore, "test")
	defer func() {
		require.NoError(t, register.Close(ctx))
	}()
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval", "return(200)"))
	defer func() {
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval")
	}()
	err := register.RegisterTask(ctx)
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		list, err := GetImportTasksFrom(ctx, client)
		return err == nil && len(list.Tasks) == 0
	}, 5*time.Second, 50*time.Millisecond)

	failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop")
	failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant")
	var list RegisterTasksList
	require.Eventually(t, func() bool {
		var err error
		list, err = GetImportTasksFrom(ctx, client)
		return err == nil && len(list.Tasks) > 0
	}, 5*time.Second, 50*time.Millisecond)
	for _, task := range list.Tasks {
		t.Log(task.MessageToUser())
		require.Equal(t, "/tidb/brie/import/restore/test", task.Key)
	}
}

func TestTaskRegisterFailedReput(t *testing.T) {
	integration.BeforeTestExternal(t)
	testEtcdCluster := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1, GRPCKeepAliveInterval: time.Second, GRPCKeepAliveTimeout: 10 * time.Second})
	defer testEtcdCluster.Terminate(t)

	ctx := context.Background()
	client := testEtcdCluster.RandClient()

	register := NewTaskRegisterWithTTL(client, 3*time.Second, RegisterRestore, "test")
	defer func() {