code.vegaprotocol.io/vega@v0.79.0/core/integration/stubs/staking_account_stub.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package stubs
    17  
    18  import (
    19  	"context"
    20  	"errors"
    21  	"fmt"
    22  	"sort"
    23  	"time"
    24  
    25  	"code.vegaprotocol.io/vega/core/types"
    26  	"code.vegaprotocol.io/vega/libs/num"
    27  )
    28  
    29  type StakingAccountStub struct {
    30  	partyToStake         map[string]*num.Uint
    31  	partyToStakeForEpoch map[int64]map[string]*num.Uint
    32  	currentEpoch         *types.Epoch
    33  }
    34  
    35  func (t *StakingAccountStub) AddEvent(ctx context.Context, evt *types.StakeLinking) {}
    36  
    37  func (t *StakingAccountStub) OnEpochEvent(ctx context.Context, epoch types.Epoch) {
    38  	if t.currentEpoch == nil || t.currentEpoch.Seq != epoch.Seq {
    39  		t.currentEpoch = &epoch
    40  		emptyT := time.Time{}
    41  		if epoch.EndTime == emptyT {
    42  			t.partyToStakeForEpoch[epoch.StartTime.UnixNano()] = map[string]*num.Uint{}
    43  			for p, s := range t.partyToStake {
    44  				t.partyToStakeForEpoch[epoch.StartTime.UnixNano()][p] = s.Clone()
    45  			}
    46  		}
    47  	}
    48  }
    49  
    50  func (t *StakingAccountStub) OnEpochRestore(_ context.Context, _ types.Epoch) {}
    51  
    52  func (t *StakingAccountStub) IncrementBalance(party string, amount *num.Uint) error {
    53  	if _, ok := t.partyToStake[party]; !ok {
    54  		t.partyToStake[party] = num.UintZero()
    55  	}
    56  	t.partyToStake[party].AddSum(amount)
    57  
    58  	return nil
    59  }
    60  
    61  func (t *StakingAccountStub) DecrementBalance(party string, amount *num.Uint) error {
    62  	if _, ok := t.partyToStake[party]; !ok {
    63  		return errors.New("party staking accoung is missing")
    64  	}
    65  	if t.partyToStake[party].LT(amount) {
    66  		return errors.New("incorrect balance for unstaking")
    67  	}
    68  	t.partyToStake[party] = t.partyToStake[party].Sub(t.partyToStake[party], amount)
    69  	t.partyToStakeForEpoch[t.currentEpoch.StartTime.UnixNano()][party] = t.partyToStake[party].Clone()
    70  	return nil
    71  }
    72  
    73  func NewStakingAccountStub() *StakingAccountStub {
    74  	return &StakingAccountStub{
    75  		partyToStake:         make(map[string]*num.Uint),
    76  		partyToStakeForEpoch: make(map[int64]map[string]*num.Uint),
    77  	}
    78  }
    79  
    80  func (t *StakingAccountStub) GetAllStakingParties() []string {
    81  	keys := make([]string, 0, len(t.partyToStake))
    82  	for k := range t.partyToStake {
    83  		keys = append(keys, k)
    84  	}
    85  	sort.Strings(keys)
    86  	return keys
    87  }
    88  
    89  func (t *StakingAccountStub) GetAvailableBalance(party string) (*num.Uint, error) {
    90  	ret, ok := t.partyToStake[party]
    91  	if !ok {
    92  		return num.UintZero(), fmt.Errorf("party not found")
    93  	}
    94  	return ret, nil
    95  }
    96  
    97  func (t *StakingAccountStub) GetAvailableBalanceInRange(party string, from, _ time.Time) (*num.Uint, error) {
    98  	partyStake, ok := t.partyToStakeForEpoch[from.UnixNano()][party]
    99  	if !ok {
   100  		return num.UintZero(), fmt.Errorf("party not found")
   101  	}
   102  	return partyStake, nil
   103  }