diff --git a/sweepbatcher/store_mock.go b/sweepbatcher/store_mock.go index 57cdd34b..815b1991 100644 --- a/sweepbatcher/store_mock.go +++ b/sweepbatcher/store_mock.go @@ -4,6 +4,7 @@ import ( "context" "errors" "sort" + "sync" "github.com/btcsuite/btcd/btcutil" "github.com/lightningnetwork/lnd/lntypes" @@ -13,6 +14,7 @@ import ( type StoreMock struct { batches map[int32]dbBatch sweeps map[lntypes.Hash]dbSweep + mu sync.Mutex } // NewStoreMock instantiates a new mock store. @@ -28,6 +30,9 @@ func NewStoreMock() *StoreMock { func (s *StoreMock) FetchUnconfirmedSweepBatches(ctx context.Context) ( []*dbBatch, error) { + s.mu.Lock() + defer s.mu.Unlock() + result := []*dbBatch{} for _, batch := range s.batches { batch := batch @@ -44,6 +49,9 @@ func (s *StoreMock) FetchUnconfirmedSweepBatches(ctx context.Context) ( func (s *StoreMock) InsertSweepBatch(ctx context.Context, batch *dbBatch) (int32, error) { + s.mu.Lock() + defer s.mu.Unlock() + var id int32 if len(s.batches) == 0 { @@ -66,12 +74,18 @@ func (s *StoreMock) DropBatch(ctx context.Context, id int32) error { func (s *StoreMock) UpdateSweepBatch(ctx context.Context, batch *dbBatch) error { + s.mu.Lock() + defer s.mu.Unlock() + s.batches[batch.ID] = *batch return nil } // ConfirmBatch confirms a batch. func (s *StoreMock) ConfirmBatch(ctx context.Context, id int32) error { + s.mu.Lock() + defer s.mu.Unlock() + batch, ok := s.batches[id] if !ok { return errors.New("batch not found") @@ -87,6 +101,9 @@ func (s *StoreMock) ConfirmBatch(ctx context.Context, id int32) error { func (s *StoreMock) FetchBatchSweeps(ctx context.Context, id int32) ([]*dbSweep, error) { + s.mu.Lock() + defer s.mu.Unlock() + result := []*dbSweep{} for _, sweep := range s.sweeps { sweep := sweep @@ -104,7 +121,11 @@ func (s *StoreMock) FetchBatchSweeps(ctx context.Context, // UpsertSweep inserts a sweep into the database, or updates an existing sweep. func (s *StoreMock) UpsertSweep(ctx context.Context, sweep *dbSweep) error { + s.mu.Lock() + defer s.mu.Unlock() + s.sweeps[sweep.SwapHash] = *sweep + return nil } @@ -112,6 +133,9 @@ func (s *StoreMock) UpsertSweep(ctx context.Context, sweep *dbSweep) error { func (s *StoreMock) GetSweepStatus(ctx context.Context, swapHash lntypes.Hash) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + sweep, ok := s.sweeps[swapHash] if !ok { return false, nil @@ -127,6 +151,9 @@ func (s *StoreMock) Close() error { // AssertSweepStored asserts that a sweep is stored. func (s *StoreMock) AssertSweepStored(id lntypes.Hash) bool { + s.mu.Lock() + defer s.mu.Unlock() + _, ok := s.sweeps[id] return ok } @@ -135,6 +162,9 @@ func (s *StoreMock) AssertSweepStored(id lntypes.Hash) bool { func (s *StoreMock) GetParentBatch(ctx context.Context, swapHash lntypes.Hash) ( *dbBatch, error) { + s.mu.Lock() + defer s.mu.Unlock() + for _, sweep := range s.sweeps { if sweep.SwapHash == swapHash { batch, ok := s.batches[sweep.BatchID] @@ -153,6 +183,9 @@ func (s *StoreMock) GetParentBatch(ctx context.Context, swapHash lntypes.Hash) ( func (s *StoreMock) TotalSweptAmount(ctx context.Context, batchID int32) ( btcutil.Amount, error) { + s.mu.Lock() + defer s.mu.Unlock() + batch, ok := s.batches[batchID] if !ok { return 0, errors.New("batch not found")