@@ -302,6 +302,29 @@ func Test_UpdateStaleData(t *testing.T) {
302302
303303}
304304
305+ func Test_DynamodbKVWithTimeout (t * testing.T ) {
306+ ddbMock := NewDynamodbClientMock ()
307+ // Backend has delay of 5s while the client timeout is 1s.
308+ ddbWithDelay := newDynamodbKVWithDelay (ddbMock , time .Second * 5 )
309+ dbWithTimeout := newDynamodbKVWithTimeout (ddbWithDelay , time .Second )
310+
311+ ctx := context .Background ()
312+ _ , _ , err := dbWithTimeout .List (ctx , dynamodbKey {primaryKey : key })
313+ require .True (t , errors .Is (err , context .DeadlineExceeded ))
314+
315+ err = dbWithTimeout .Delete (ctx , dynamodbKey {primaryKey : key })
316+ require .True (t , errors .Is (err , context .DeadlineExceeded ))
317+
318+ _ , _ , err = dbWithTimeout .Query (ctx , dynamodbKey {primaryKey : key }, true )
319+ require .True (t , errors .Is (err , context .DeadlineExceeded ))
320+
321+ err = dbWithTimeout .Put (ctx , dynamodbKey {primaryKey : key }, []byte {})
322+ require .True (t , errors .Is (err , context .DeadlineExceeded ))
323+
324+ err = dbWithTimeout .Batch (ctx , nil , nil )
325+ require .True (t , errors .Is (err , context .DeadlineExceeded ))
326+ }
327+
305328// NewClientMock makes a new local dynamodb client.
306329func NewClientMock (ddbClient dynamoDbClient , cc codec.Codec , logger log.Logger , registerer prometheus.Registerer , time time.Duration , config backoff.Config ) * Client {
307330 return & Client {
@@ -429,3 +452,57 @@ func (m *DescMock) FindDifference(that codec.MultiKey) (interface{}, []string, e
429452 }
430453 return args .Get (0 ), args .Get (1 ).([]string ), err
431454}
455+
456+ type dynamodbKVWithDelayAndContextCheck struct {
457+ ddbClient dynamoDbClient
458+ delay time.Duration
459+ }
460+
461+ func newDynamodbKVWithDelay (client dynamoDbClient , delay time.Duration ) * dynamodbKVWithDelayAndContextCheck {
462+ return & dynamodbKVWithDelayAndContextCheck {ddbClient : client , delay : delay }
463+ }
464+
465+ func (d * dynamodbKVWithDelayAndContextCheck ) List (ctx context.Context , key dynamodbKey ) ([]string , float64 , error ) {
466+ select {
467+ case <- ctx .Done ():
468+ return nil , 0 , ctx .Err ()
469+ case <- time .After (d .delay ):
470+ return d .ddbClient .List (ctx , key )
471+ }
472+ }
473+
474+ func (d * dynamodbKVWithDelayAndContextCheck ) Query (ctx context.Context , key dynamodbKey , isPrefix bool ) (map [string ][]byte , float64 , error ) {
475+ select {
476+ case <- ctx .Done ():
477+ return nil , 0 , ctx .Err ()
478+ case <- time .After (d .delay ):
479+ return d .ddbClient .Query (ctx , key , isPrefix )
480+ }
481+ }
482+
483+ func (d * dynamodbKVWithDelayAndContextCheck ) Delete (ctx context.Context , key dynamodbKey ) error {
484+ select {
485+ case <- ctx .Done ():
486+ return ctx .Err ()
487+ case <- time .After (d .delay ):
488+ return d .ddbClient .Delete (ctx , key )
489+ }
490+ }
491+
492+ func (d * dynamodbKVWithDelayAndContextCheck ) Put (ctx context.Context , key dynamodbKey , data []byte ) error {
493+ select {
494+ case <- ctx .Done ():
495+ return ctx .Err ()
496+ case <- time .After (d .delay ):
497+ return d .ddbClient .Put (ctx , key , data )
498+ }
499+ }
500+
501+ func (d * dynamodbKVWithDelayAndContextCheck ) Batch (ctx context.Context , put map [dynamodbKey ][]byte , delete []dynamodbKey ) error {
502+ select {
503+ case <- ctx .Done ():
504+ return ctx .Err ()
505+ case <- time .After (d .delay ):
506+ return d .ddbClient .Batch (ctx , put , delete )
507+ }
508+ }
0 commit comments