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"
"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

View File

@ -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)
}

View File

@ -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"