mirror of https://github.com/tikv/client-go.git
Fix the test
Signed-off-by: JmPotato <ghzpotato@gmail.com>
This commit is contained in:
parent
08694efec4
commit
8553f26149
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
Loading…
Reference in New Issue