br: fix register_test failed error for flaky test (#46991)

close #46848
pingcap · Feb 7, 2024 · dee790b · dee790b
1 parent 9553bc6
commit dee790b
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/br/pkg/utils/register_test.go b/br/pkg/utils/register_test.go
@@ -15,6 +15,7 @@ package utils
 
 import (
 	"context"
+	"strings"
 	"testing"
 	"time"
 
@@ -115,8 +116,12 @@ func TestTaskRegisterFailedGrant(t *testing.T) {
 		t.Log(task.MessageToUser())
 		require.Equal(t, "/tidb/brie/import/restore/test", task.Key)
 	}
-
-	require.NoError(t, register.Close(ctx))
+	require.True(t, len(list.Tasks) > 0)
+	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)
+	}
 }
 
 func TestTaskRegisterFailedReput(t *testing.T) {
@@ -153,6 +158,10 @@ func TestTaskRegisterFailedReput(t *testing.T) {
 		t.Log(task.MessageToUser())
 		require.Equal(t, "/tidb/brie/import/restore/test", task.Key)
 	}
-
-	require.NoError(t, register.Close(ctx))
+	require.True(t, len(list.Tasks) > 0)
+	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)
+	}
 }