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"
|
||||
"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/interceptor"
|
||||
"github.com/tikv/pd/pkg/mcs/resource_manager/client"
|
||||
)
|
||||
|
||||
func init() {
|
||||
resourceControlSwitch.Store(false)
|
||||
resourceControlInterceptor = nil
|
||||
}
|
||||
|
||||
var _ Client = interceptedClient{}
|
||||
|
||||
type interceptedClient struct {
|
||||
|
|
@ -101,7 +106,7 @@ func buildResourceControlInterceptor(
|
|||
return nil
|
||||
}
|
||||
// Make the request info.
|
||||
reqInfo := resource_control.MakeRequestInfo(req)
|
||||
reqInfo := resourcecontrol.MakeRequestInfo(req)
|
||||
// Build the interceptor.
|
||||
return func(next interceptor.RPCInterceptorFunc) interceptor.RPCInterceptorFunc {
|
||||
return func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
|
||||
|
|
@ -111,7 +116,7 @@ func buildResourceControlInterceptor(
|
|||
}
|
||||
resp, err := next(target, req)
|
||||
if resp != nil {
|
||||
respInfo := resource_control.MakeResponseInfo(resp)
|
||||
respInfo := resourcecontrol.MakeResponseInfo(resp)
|
||||
resourceControlInterceptor.OnResponse(ctx, resourceGroupName, reqInfo, respInfo)
|
||||
}
|
||||
return resp, err
|
||||
|
|
|
|||
|
|
@ -47,6 +47,6 @@ func TestInterceptedClient(t *testing.T) {
|
|||
return next(target, req)
|
||||
}
|
||||
})
|
||||
_, _ = client.SendRequest(ctx, "", nil, 0)
|
||||
_, _ = client.SendRequest(ctx, "", &tikvrpc.Request{}, 0)
|
||||
assert.True(t, executed)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package resource_control
|
||||
package resourcecontrol
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
Loading…
Reference in New Issue