Fix the test

Signed-off-by: JmPotato <ghzpotato@gmail.com>
This commit is contained in:
JmPotato 2023-01-19 16:55:26 +08:00
parent 08694efec4
commit 8553f26149
3 changed files with 10 additions and 5 deletions

View File

@ -19,12 +19,17 @@ import (
"sync/atomic" "sync/atomic"
"time" "time"
"github.com/tikv/client-go/v2/internal/resource_control" "github.com/tikv/client-go/v2/internal/resourcecontrol"
"github.com/tikv/client-go/v2/tikvrpc" "github.com/tikv/client-go/v2/tikvrpc"
"github.com/tikv/client-go/v2/tikvrpc/interceptor" "github.com/tikv/client-go/v2/tikvrpc/interceptor"
"github.com/tikv/pd/pkg/mcs/resource_manager/client" "github.com/tikv/pd/pkg/mcs/resource_manager/client"
) )
func init() {
resourceControlSwitch.Store(false)
resourceControlInterceptor = nil
}
var _ Client = interceptedClient{} var _ Client = interceptedClient{}
type interceptedClient struct { type interceptedClient struct {
@ -101,7 +106,7 @@ func buildResourceControlInterceptor(
return nil return nil
} }
// Make the request info. // Make the request info.
reqInfo := resource_control.MakeRequestInfo(req) reqInfo := resourcecontrol.MakeRequestInfo(req)
// Build the interceptor. // Build the interceptor.
return func(next interceptor.RPCInterceptorFunc) interceptor.RPCInterceptorFunc { return func(next interceptor.RPCInterceptorFunc) interceptor.RPCInterceptorFunc {
return func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) { return func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
@ -111,7 +116,7 @@ func buildResourceControlInterceptor(
} }
resp, err := next(target, req) resp, err := next(target, req)
if resp != nil { if resp != nil {
respInfo := resource_control.MakeResponseInfo(resp) respInfo := resourcecontrol.MakeResponseInfo(resp)
resourceControlInterceptor.OnResponse(ctx, resourceGroupName, reqInfo, respInfo) resourceControlInterceptor.OnResponse(ctx, resourceGroupName, reqInfo, respInfo)
} }
return resp, err return resp, err

View File

@ -47,6 +47,6 @@ func TestInterceptedClient(t *testing.T) {
return next(target, req) return next(target, req)
} }
}) })
_, _ = client.SendRequest(ctx, "", nil, 0) _, _ = client.SendRequest(ctx, "", &tikvrpc.Request{}, 0)
assert.True(t, executed) assert.True(t, executed)
} }

View File

@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
package resource_control package resourcecontrol
import ( import (
"reflect" "reflect"