expose GetRestyClient

dtmcli/types.go

@@ -55,14 +55,9 @@ func SetBarrierTableName(tablename string) {
 	dtmimp.BarrierTableName = tablename
 }
 
-// OnBeforeRequest add before request middleware
-func OnBeforeRequest(middleware func(c *resty.Client, r *resty.Request) error) {
-	dtmimp.RestyClient.OnBeforeRequest(middleware)
-}
-
-// OnAfterResponse add after request middleware
-func OnAfterResponse(middleware func(c *resty.Client, resp *resty.Response) error) {
-	dtmimp.RestyClient.OnAfterResponse(middleware)
+// GetRestyClient get the resty.Client for http request
+func GetRestyClient() *resty.Client {
+	return dtmimp.RestyClient
 }
 
 // SetPassthroughHeaders experimental.

test/main_test.go

@@ -37,8 +37,8 @@ func TestMain(m *testing.M) {
 	conf.UpdateBranchSync = 1
 
 	dtmgrpc.AddUnaryInterceptor(busi.SetGrpcHeaderForHeadersYes)
-	dtmcli.OnBeforeRequest(busi.SetHttpHeaderForHeadersYes)
-	dtmcli.OnAfterResponse(func(c *resty.Client, resp *resty.Response) error { return nil })
+	dtmcli.GetRestyClient().OnBeforeRequest(busi.SetHttpHeaderForHeadersYes)
+	dtmcli.GetRestyClient().OnAfterResponse(func(c *resty.Client, resp *resty.Response) error { return nil })
 
 	tenv := os.Getenv("TEST_STORE")
 	if tenv == "boltdb" {