Repository: pingcap/tidb
File: pkg/executor/internal/querywatch/query_watch_test.go (Go)

Code before:
	require.Equal(t, 1, len(r.Rows()))
	rs, err = tk.Exec("SET @rg=rg2")
	require.NoError(t, err)
	require.Nil(t, rs)
	rs, err = tk.Exec("query watch remove resource group @rg")
	require.NoError(t, err)
	require.Nil(t, rs)
	r = tk.MustQuery("select * from information_schema.runaway_watches where resource_group_name = 'rg2'")
	require.Equal(t, 0, len(r.Rows()))
	// test remove by id
	rs, err = tk.Exec("query watch remove 1")
	require.NoError(t, err)
	require.Nil(t, rs)
	require.Eventually(t, func() bool {
		err := tk.ExecToErr("select * from test.t1")
		if err == nil {
			return false
		}
		originErr := errors.Cause(err)
		tErr, ok := originErr.(*terror.Error)
		if !ok {
			return false
		}
		sqlErr := terror.ToSQLError(tErr)
		return int(sqlErr.Code) == mysql.ErrResourceGroupQueryRunawayQuarantine
	}, 10*time.Second, 200*time.Millisecond)
}

func TestQueryWatchIssue56897(t *testing.T) {
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/resourcegroup/runaway/FastRunawayGC", `return(1)`))
	defer func() {
		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/resourcegroup/runaway/FastRunawayGC"))
	}()
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)
	tk.MustExec("use test")
	tk.MustQuery("QUERY WATCH ADD ACTION KILL SQL TEXT SIMILAR TO 'use test';").Check((testkit.Rows("1")))
	time.Sleep(1 * time.Second)
	_, err := tk.Exec("use test")
	require.Nil(t, err)
	_, err = tk.Exec("use mysql")
	require.Nil(t, err)
}


Review comment (suggestion):
The retry interval is hard-coded as `200*time.Millisecond` here even though `tryInterval` is already defined earlier in the test. Reusing that variable would keep retry behavior consistent and avoid duplication (or introduce a new local `waitFor` variable if you intentionally want a different timeout than `maxWaitDuration`).
```suggestion
	}, 10*time.Second, tryInterval)
```

Code after:
	require.Equal(t, 1, len(r.Rows()))
	rs, err = tk.Exec("SET @rg=rg2")
	require.NoError(t, err)
	require.Nil(t, rs)
	rs, err = tk.Exec("query watch remove resource group @rg")
	require.NoError(t, err)
	require.Nil(t, rs)
	r = tk.MustQuery("select * from information_schema.runaway_watches where resource_group_name = 'rg2'")
	require.Equal(t, 0, len(r.Rows()))
	// test remove by id
	rs, err = tk.Exec("query watch remove 1")
	require.NoError(t, err)
	require.Nil(t, rs)
	var lastObservedErr any
	require.Eventuallyf(t, func() bool {
		err := tk.ExecToErr("select * from test.t1")
		if err == nil {
			lastObservedErr = nil
			return false
		}
		originErr := errors.Cause(err)
		tErr, ok := originErr.(*terror.Error)
		if !ok {
			lastObservedErr = originErr
			return false
		}
		sqlErr := terror.ToSQLError(tErr)
		if int(sqlErr.Code) != mysql.ErrResourceGroupQueryRunawayQuarantine {
			lastObservedErr = sqlErr
			return false
		}
		return true
	}, 10*time.Second, tryInterval, "expected quarantine error (%d), last observed: %v", mysql.ErrResourceGroupQueryRunawayQuarantine, lastObservedErr)
}

func TestQueryWatchIssue56897(t *testing.T) {
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/resourcegroup/runaway/FastRunawayGC", `return(1)`))
	defer func() {
		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/resourcegroup/runaway/FastRunawayGC"))
	}()
	store := testkit.CreateMockStore(t)
	tk := testkit.NewTestKit(t, store)
	tk.MustExec("use test")
	tk.MustQuery("QUERY WATCH ADD ACTION KILL SQL TEXT SIMILAR TO 'use test';").Check((testkit.Rows("1")))
	time.Sleep(1 * time.Second)
	_, err := tk.Exec("use test")
	require.Nil(t, err)
	_, err = tk.Exec("use mysql")
	require.Nil(t, err)
}