github.com/openfga/openfga@v1.5.4-rc1/internal/mocks/mock_slow_storage.go (about)

     1  package mocks
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	openfgav1 "github.com/openfga/api/proto/openfga/v1"
     8  
     9  	"github.com/openfga/openfga/pkg/storage"
    10  )
    11  
    12  // slowDataStorage is a proxy to the actual ds except the Reads are slow time by the readTuplesDelay
    13  // This allows simulating list objection condition that times out.
    14  type slowDataStorage struct {
    15  	readTuplesDelay time.Duration
    16  	storage.OpenFGADatastore
    17  }
    18  
    19  // NewMockSlowDataStorage returns a wrapper of a datastore that adds artificial delays into the reads of tuples.
    20  func NewMockSlowDataStorage(ds storage.OpenFGADatastore, readTuplesDelay time.Duration) storage.OpenFGADatastore {
    21  	return &slowDataStorage{
    22  		readTuplesDelay:  readTuplesDelay,
    23  		OpenFGADatastore: ds,
    24  	}
    25  }
    26  
    27  func (m *slowDataStorage) Close() {}
    28  
    29  func (m *slowDataStorage) Read(ctx context.Context, store string, key *openfgav1.TupleKey) (storage.TupleIterator, error) {
    30  	time.Sleep(m.readTuplesDelay)
    31  	return m.OpenFGADatastore.Read(ctx, store, key)
    32  }
    33  
    34  func (m *slowDataStorage) ReadPage(ctx context.Context, store string, key *openfgav1.TupleKey, paginationOptions storage.PaginationOptions) ([]*openfgav1.Tuple, []byte, error) {
    35  	time.Sleep(m.readTuplesDelay)
    36  	return m.OpenFGADatastore.ReadPage(ctx, store, key, paginationOptions)
    37  }
    38  
    39  func (m *slowDataStorage) ReadUserTuple(ctx context.Context, store string, key *openfgav1.TupleKey) (*openfgav1.Tuple, error) {
    40  	time.Sleep(m.readTuplesDelay)
    41  	return m.OpenFGADatastore.ReadUserTuple(ctx, store, key)
    42  }
    43  
    44  func (m *slowDataStorage) ReadUsersetTuples(ctx context.Context, store string, filter storage.ReadUsersetTuplesFilter) (storage.TupleIterator, error) {
    45  	time.Sleep(m.readTuplesDelay)
    46  	return m.OpenFGADatastore.ReadUsersetTuples(ctx, store, filter)
    47  }
    48  
    49  func (m *slowDataStorage) ReadStartingWithUser(
    50  	ctx context.Context,
    51  	store string,
    52  	filter storage.ReadStartingWithUserFilter,
    53  ) (storage.TupleIterator, error) {
    54  	time.Sleep(m.readTuplesDelay)
    55  	return m.OpenFGADatastore.ReadStartingWithUser(ctx, store, filter)
    56  }