diff options
-rw-r--r-- | internal/gitaly/transaction/mock.go | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/internal/gitaly/transaction/mock.go b/internal/gitaly/transaction/mock.go index 101d4b291..0d12b9f91 100644 --- a/internal/gitaly/transaction/mock.go +++ b/internal/gitaly/transaction/mock.go @@ -3,6 +3,7 @@ package transaction import ( "context" "errors" + "sync" "gitlab.com/gitlab-org/gitaly/v14/internal/transaction/txinfo" "gitlab.com/gitlab-org/gitaly/v14/internal/transaction/voting" @@ -29,3 +30,44 @@ func (m *MockManager) Stop(ctx context.Context, tx txinfo.Transaction) error { } return m.StopFn(ctx, tx) } + +// TrackingManager is a transaction manager which tracks all votes. Voting functions never return +// an error. +type TrackingManager struct { + MockManager + + votesLock sync.Mutex + votes []voting.Vote +} + +// NewTrackingManager creates a new TrackingManager which is ready for use. +func NewTrackingManager() *TrackingManager { + manager := &TrackingManager{} + + manager.VoteFn = func(_ context.Context, _ txinfo.Transaction, vote voting.Vote) error { + manager.votesLock.Lock() + defer manager.votesLock.Unlock() + manager.votes = append(manager.votes, vote) + return nil + } + + return manager +} + +// Votes returns a copy of all votes which have been cast. +func (m *TrackingManager) Votes() []voting.Vote { + m.votesLock.Lock() + defer m.votesLock.Unlock() + + votes := make([]voting.Vote, len(m.votes)) + copy(votes, m.votes) + + return votes +} + +// Reset resets all votes which have been recorded up to this point. +func (m *TrackingManager) Reset() { + m.votesLock.Lock() + defer m.votesLock.Unlock() + m.votes = nil +} |