loop/sweepbatcher/store_mock.go

package sweepbatcher

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"sync"

	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/wire"
)

// StoreMock implements a mock client swap store.
type StoreMock struct {
	batches map[int32]dbBatch
	sweeps  map[wire.OutPoint]dbSweep
	mu      sync.Mutex
	sweepID int32
	batchID int32
}

// NewStoreMock instantiates a new mock store.
func NewStoreMock() *StoreMock {
	return &StoreMock{
		batches: make(map[int32]dbBatch),
		sweeps:  make(map[wire.OutPoint]dbSweep),
	}
}

// FetchUnconfirmedSweepBatches fetches all the loop out sweep batches from the
// database that are not in a confirmed state.
func (s *StoreMock) FetchUnconfirmedSweepBatches(ctx context.Context) (
	[]*dbBatch, error) {

	s.mu.Lock()
	defer s.mu.Unlock()

	result := []*dbBatch{}
	for _, batch := range s.batches {
		if !batch.Confirmed {
			result = append(result, &batch)
		}
	}

	return result, nil
}

// InsertSweepBatch inserts a batch into the database, returning the id of the
// inserted batch.
func (s *StoreMock) InsertSweepBatch(ctx context.Context,
	batch *dbBatch) (int32, error) {

	s.mu.Lock()
	defer s.mu.Unlock()

	id := s.batchID
	s.batchID++

	s.batches[id] = *batch
	return id, nil
}

// CancelBatch drops a batch from the database.
func (s *StoreMock) CancelBatch(ctx context.Context, id int32) error {
	delete(s.batches, id)
	return nil
}

// UpdateSweepBatch updates a batch in the database.
func (s *StoreMock) UpdateSweepBatch(ctx context.Context,
	batch *dbBatch) error {

	s.mu.Lock()
	defer s.mu.Unlock()

	s.batches[batch.ID] = *batch
	return nil
}

// ConfirmBatchWithSweeps updates the batch and the provided sweeps atomically.
func (s *StoreMock) ConfirmBatchWithSweeps(ctx context.Context,
	batch *dbBatch, sweeps []*dbSweep) error {

	s.mu.Lock()
	defer s.mu.Unlock()

	s.batches[batch.ID] = *batch

	for _, sweep := range sweeps {
		sweepCopy := *sweep

		old, exists := s.sweeps[sweep.Outpoint]
		if !exists {
			return fmt.Errorf("confirming unknown sweep %v",
				sweep.Outpoint)
		}

		sweepCopy.ID = old.ID

		s.sweeps[sweep.Outpoint] = sweepCopy
	}

	return nil
}

// FetchBatchSweeps fetches all the sweeps that belong to a batch.
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 {
		if sweep.BatchID == id {
			result = append(result, &sweep)
		}
	}

	sort.Slice(result, func(i, j int) bool {
		return result[i].ID < result[j].ID
	})

	return result, nil
}

// 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()

	sweepCopy := *sweep

	if old, exists := s.sweeps[sweep.Outpoint]; exists {
		// Preserve existing sweep ID.
		sweepCopy.ID = old.ID
	} else {
		// Assign fresh sweep ID.
		sweepCopy.ID = s.sweepID
		s.sweepID++
	}

	s.sweeps[sweep.Outpoint] = sweepCopy

	return nil
}

// GetSweepStatus returns the status of a sweep.
func (s *StoreMock) GetSweepStatus(ctx context.Context,
	outpoint wire.OutPoint) (bool, error) {

	s.mu.Lock()
	defer s.mu.Unlock()

	sweep, ok := s.sweeps[outpoint]
	if !ok {
		return false, nil
	}

	return sweep.Completed, nil
}

// Close closes the store.
func (s *StoreMock) Close() error {
	return nil
}

// AssertSweepStored asserts that a sweep is stored.
func (s *StoreMock) AssertSweepStored(outpoint wire.OutPoint) bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, ok := s.sweeps[outpoint]
	return ok
}

// GetParentBatch returns the parent batch of a swap.
func (s *StoreMock) GetParentBatch(ctx context.Context,
	outpoint wire.OutPoint) (*dbBatch, error) {

	s.mu.Lock()
	defer s.mu.Unlock()

	for _, sweep := range s.sweeps {
		if sweep.Outpoint == outpoint {
			batch, ok := s.batches[sweep.BatchID]
			if !ok {
				return nil, errors.New("batch not found")
			}
			return &batch, nil
		}
	}

	return nil, errors.New("batch not found")
}

// TotalSweptAmount returns the total amount of BTC that has been swept from a
// batch.
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")
	}

	if !batch.Confirmed {
		return 0, nil
	}

	var total btcutil.Amount
	for _, sweep := range s.sweeps {
		if sweep.BatchID == batchID {
			total += sweep.Amount
		}
	}

	return total, nil
}