7
7
8
8
use crate :: {
9
9
pd:: { PdClient , PdRpcClient , RetryClient } ,
10
- region:: { Region , RegionId } ,
11
- store:: Store ,
10
+ region:: { RegionId , RegionWithLeader } ,
11
+ store:: RegionStore ,
12
12
Config , Error , Key , Result , Timestamp ,
13
13
} ;
14
14
use async_trait:: async_trait;
15
15
use derive_new:: new;
16
+ use slog:: { Drain , Logger } ;
16
17
use std:: { any:: Any , sync:: Arc } ;
17
18
use tikv_client_proto:: metapb;
18
19
use tikv_client_store:: { KvClient , KvConnect , Request } ;
@@ -21,8 +22,16 @@ use tikv_client_store::{KvClient, KvConnect, Request};
21
22
/// client can be tested without doing any RPC calls.
22
23
pub async fn pd_rpc_client ( ) -> PdRpcClient < MockKvConnect , MockCluster > {
23
24
let config = Config :: default ( ) ;
25
+ let plain = slog_term:: PlainSyncDecorator :: new ( std:: io:: stdout ( ) ) ;
26
+ let logger = Logger :: root (
27
+ slog_term:: FullFormat :: new ( plain)
28
+ . build ( )
29
+ . filter_level ( slog:: Level :: Info )
30
+ . fuse ( ) ,
31
+ o ! ( ) ,
32
+ ) ;
24
33
PdRpcClient :: new (
25
- & config,
34
+ config. clone ( ) ,
26
35
|_, _| MockKvConnect ,
27
36
|e, sm| {
28
37
futures:: future:: ok ( RetryClient :: new_with_cluster (
@@ -33,11 +42,13 @@ pub async fn pd_rpc_client() -> PdRpcClient<MockKvConnect, MockCluster> {
33
42
) )
34
43
} ,
35
44
false ,
45
+ logger,
36
46
)
37
47
. await
38
48
. unwrap ( )
39
49
}
40
50
51
+ #[ allow( clippy:: type_complexity) ]
41
52
#[ derive( new, Default , Clone ) ]
42
53
pub struct MockKvClient {
43
54
pub addr : String ,
@@ -93,27 +104,31 @@ impl MockPdClient {
93
104
}
94
105
}
95
106
96
- pub fn region1 ( ) -> Region {
97
- let mut region = Region :: default ( ) ;
107
+ pub fn region1 ( ) -> RegionWithLeader {
108
+ let mut region = RegionWithLeader :: default ( ) ;
98
109
region. region . id = 1 ;
99
110
region. region . set_start_key ( vec ! [ 0 ] ) ;
100
111
region. region . set_end_key ( vec ! [ 10 ] ) ;
101
112
102
- let mut leader = metapb:: Peer :: default ( ) ;
103
- leader. store_id = 41 ;
113
+ let leader = metapb:: Peer {
114
+ store_id : 41 ,
115
+ ..Default :: default ( )
116
+ } ;
104
117
region. leader = Some ( leader) ;
105
118
106
119
region
107
120
}
108
121
109
- pub fn region2 ( ) -> Region {
110
- let mut region = Region :: default ( ) ;
122
+ pub fn region2 ( ) -> RegionWithLeader {
123
+ let mut region = RegionWithLeader :: default ( ) ;
111
124
region. region . id = 2 ;
112
125
region. region . set_start_key ( vec ! [ 10 ] ) ;
113
126
region. region . set_end_key ( vec ! [ 250 , 250 ] ) ;
114
127
115
- let mut leader = metapb:: Peer :: default ( ) ;
116
- leader. store_id = 42 ;
128
+ let leader = metapb:: Peer {
129
+ store_id : 42 ,
130
+ ..Default :: default ( )
131
+ } ;
117
132
region. leader = Some ( leader) ;
118
133
119
134
region
@@ -124,11 +139,11 @@ impl MockPdClient {
124
139
impl PdClient for MockPdClient {
125
140
type KvClient = MockKvClient ;
126
141
127
- async fn map_region_to_store ( self : Arc < Self > , region : Region ) -> Result < Store > {
128
- Ok ( Store :: new ( region, Arc :: new ( self . client . clone ( ) ) ) )
142
+ async fn map_region_to_store ( self : Arc < Self > , region : RegionWithLeader ) -> Result < RegionStore > {
143
+ Ok ( RegionStore :: new ( region, Arc :: new ( self . client . clone ( ) ) ) )
129
144
}
130
145
131
- async fn region_for_key ( & self , key : & Key ) -> Result < Region > {
146
+ async fn region_for_key ( & self , key : & Key ) -> Result < RegionWithLeader > {
132
147
let bytes: & [ _ ] = key. into ( ) ;
133
148
let region = if bytes. is_empty ( ) || bytes[ 0 ] < 10 {
134
149
Self :: region1 ( )
@@ -139,11 +154,11 @@ impl PdClient for MockPdClient {
139
154
Ok ( region)
140
155
}
141
156
142
- async fn region_for_id ( & self , id : RegionId ) -> Result < Region > {
157
+ async fn region_for_id ( & self , id : RegionId ) -> Result < RegionWithLeader > {
143
158
match id {
144
159
1 => Ok ( Self :: region1 ( ) ) ,
145
160
2 => Ok ( Self :: region2 ( ) ) ,
146
- _ => Err ( Error :: RegionNotFound { region_id : id } ) ,
161
+ _ => Err ( Error :: RegionNotFoundInResponse { region_id : id } ) ,
147
162
}
148
163
}
149
164
@@ -154,11 +169,21 @@ impl PdClient for MockPdClient {
154
169
async fn update_safepoint ( self : Arc < Self > , _safepoint : u64 ) -> Result < bool > {
155
170
unimplemented ! ( )
156
171
}
172
+
173
+ async fn update_leader (
174
+ & self ,
175
+ _ver_id : crate :: region:: RegionVerId ,
176
+ _leader : metapb:: Peer ,
177
+ ) -> Result < ( ) > {
178
+ todo ! ( )
179
+ }
180
+
181
+ async fn invalidate_region_cache ( & self , _ver_id : crate :: region:: RegionVerId ) { }
157
182
}
158
183
159
- pub fn mock_store ( ) -> Store {
160
- Store {
161
- region : Region :: default ( ) ,
184
+ pub fn mock_store ( ) -> RegionStore {
185
+ RegionStore {
186
+ region_with_leader : RegionWithLeader :: default ( ) ,
162
187
client : Arc :: new ( MockKvClient :: new ( "foo" . to_owned ( ) , None ) ) ,
163
188
}
164
189
}
0 commit comments