github.com/KinWaiYuen/client-go/v2@v2.5.4/rawkv/test_prob.go (about)

     1  // Copyright 2021 TiKV Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package rawkv
    16  
    17  import (
    18  	"github.com/KinWaiYuen/client-go/v2/internal/client"
    19  	"github.com/KinWaiYuen/client-go/v2/internal/locate"
    20  	pd "github.com/tikv/pd/client"
    21  )
    22  
    23  // ClientProbe wraps RawKVClient and exposes internal states for testing purpose.
    24  type ClientProbe struct {
    25  	*Client
    26  }
    27  
    28  // GetRegionCache returns the internal region cache container.
    29  func (c ClientProbe) GetRegionCache() *locate.RegionCache {
    30  	return c.regionCache
    31  }
    32  
    33  // SetRegionCache resets the internal region cache container.
    34  func (c ClientProbe) SetRegionCache(regionCache *locate.RegionCache) {
    35  	c.regionCache = regionCache
    36  }
    37  
    38  // SetPDClient resets the interval PD client.
    39  func (c ClientProbe) SetPDClient(client pd.Client) {
    40  	c.pdClient = client
    41  }
    42  
    43  // SetRPCClient resets the internal RPC client.
    44  func (c ClientProbe) SetRPCClient(client client.Client) {
    45  	c.rpcClient = client
    46  }
    47  
    48  // ConfigProbe exposes configurations and global variables for testing purpose.
    49  type ConfigProbe struct{}
    50  
    51  // GetRawBatchPutSize returns the raw batch put size config.
    52  func (c ConfigProbe) GetRawBatchPutSize() int {
    53  	return rawBatchPutSize
    54  }