diff --git a/tests/server/tso/allocator_test.go b/tests/server/tso/allocator_test.go
index 8b1ad3dc7d1..41f544729c2 100644
--- a/tests/server/tso/allocator_test.go
+++ b/tests/server/tso/allocator_test.go
@@ -215,7 +215,7 @@ func testTSOSuffix(re *require.Assertions, cluster *tests.TestCluster, am *tso.A
 	re.NoError(err)
 	var tso pdpb.Timestamp
 	testutil.Eventually(re, func() bool {
-		tso, err = allocator.GenerateTSO(1)
+		tso, err = allocator.GenerateTSO(context.Background(), 1)
 		re.NoError(err)
 		return tso.GetPhysical() != 0
 	})