@@ -66,15 +66,17 @@ mod test {
6666 use crate :: {
6767 mock:: { MockKvClient , MockPdClient } ,
6868 store:: store_stream_for_keys,
69+ transaction:: lowering:: new_commit_request,
6970 Error , Key , Result ,
7071 } ;
7172 use futures:: executor;
7273 use grpcio:: CallOption ;
7374 use std:: {
7475 any:: Any ,
76+ iter,
7577 sync:: { Arc , Mutex } ,
7678 } ;
77- use tikv_client_proto:: { kvrpcpb, tikvpb:: TikvClient } ;
79+ use tikv_client_proto:: { kvrpcpb, pdpb :: Timestamp , tikvpb:: TikvClient } ;
7880 use tikv_client_store:: HasRegionError ;
7981
8082 #[ test]
@@ -174,4 +176,47 @@ mod test {
174176 // Original call plus the 3 retries
175177 assert_eq ! ( * invoking_count. lock( ) . unwrap( ) , 4 ) ;
176178 }
179+
180+ #[ test]
181+ fn test_extract_error ( ) {
182+ let pd_client = Arc :: new ( MockPdClient :: new ( MockKvClient :: with_dispatch_hook (
183+ |_: & dyn Any | {
184+ Ok ( Box :: new ( kvrpcpb:: CommitResponse {
185+ region_error : None ,
186+ error : Some ( kvrpcpb:: KeyError {
187+ locked : None ,
188+ retryable : String :: new ( ) ,
189+ abort : String :: new ( ) ,
190+ conflict : None ,
191+ already_exist : None ,
192+ deadlock : None ,
193+ commit_ts_expired : None ,
194+ txn_not_found : None ,
195+ commit_ts_too_large : None ,
196+ } ) ,
197+ commit_version : 0 ,
198+ } ) as Box < dyn Any > )
199+ } ,
200+ ) ) ) ;
201+
202+ let key: Key = "key" . to_owned ( ) . into ( ) ;
203+ let req = new_commit_request ( iter:: once ( key) , Timestamp :: default ( ) , Timestamp :: default ( ) ) ;
204+
205+ // does not extract error
206+ let plan = crate :: request:: PlanBuilder :: new ( pd_client. clone ( ) , req. clone ( ) )
207+ . resolve_lock ( OPTIMISTIC_BACKOFF )
208+ . multi_region ( )
209+ . retry_region ( OPTIMISTIC_BACKOFF )
210+ . plan ( ) ;
211+ assert ! ( executor:: block_on( async { plan. execute( ) . await } ) . is_ok( ) ) ;
212+
213+ // extract error
214+ let plan = crate :: request:: PlanBuilder :: new ( pd_client. clone ( ) , req)
215+ . resolve_lock ( OPTIMISTIC_BACKOFF )
216+ . multi_region ( )
217+ . retry_region ( OPTIMISTIC_BACKOFF )
218+ . extract_error ( )
219+ . plan ( ) ;
220+ assert ! ( executor:: block_on( async { plan. execute( ) . await } ) . is_err( ) ) ;
221+ }
177222}
0 commit comments