github.com/ethersphere/bee/v2@v2.2.0/pkg/storageincentives/staking/mock/contract.go (about)

     1  // Copyright 2022 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package mock
     6  
     7  import (
     8  	"context"
     9  	"math/big"
    10  
    11  	"github.com/ethereum/go-ethereum/common"
    12  	"github.com/ethersphere/bee/v2/pkg/storageincentives/staking"
    13  )
    14  
    15  type stakingContractMock struct {
    16  	depositStake     func(ctx context.Context, stakedAmount *big.Int) (common.Hash, error)
    17  	getStake         func(ctx context.Context) (*big.Int, error)
    18  	withdrawAllStake func(ctx context.Context) (common.Hash, error)
    19  	migrateStake     func(ctx context.Context) (common.Hash, error)
    20  	isFrozen         func(ctx context.Context, block uint64) (bool, error)
    21  }
    22  
    23  func (s *stakingContractMock) DepositStake(ctx context.Context, stakedAmount *big.Int) (common.Hash, error) {
    24  	return s.depositStake(ctx, stakedAmount)
    25  }
    26  
    27  func (s *stakingContractMock) ChangeStakeOverlay(_ context.Context, h common.Hash) (common.Hash, error) {
    28  	return h, nil
    29  }
    30  
    31  func (s *stakingContractMock) GetPotentialStake(ctx context.Context) (*big.Int, error) {
    32  	return s.getStake(ctx)
    33  }
    34  
    35  func (s *stakingContractMock) GetWithdrawableStake(ctx context.Context) (*big.Int, error) {
    36  	return s.getStake(ctx)
    37  }
    38  
    39  func (s *stakingContractMock) WithdrawStake(ctx context.Context) (common.Hash, error) {
    40  	return s.withdrawAllStake(ctx)
    41  }
    42  
    43  func (s *stakingContractMock) MigrateStake(ctx context.Context) (common.Hash, error) {
    44  	return s.migrateStake(ctx)
    45  }
    46  
    47  func (s *stakingContractMock) IsOverlayFrozen(ctx context.Context, block uint64) (bool, error) {
    48  	return s.isFrozen(ctx, block)
    49  }
    50  
    51  // Option is a an option passed to New
    52  type Option func(mock *stakingContractMock)
    53  
    54  // New creates a new mock BatchStore.
    55  func New(opts ...Option) staking.Contract {
    56  	bs := &stakingContractMock{}
    57  
    58  	for _, o := range opts {
    59  		o(bs)
    60  	}
    61  
    62  	return bs
    63  }
    64  
    65  func WithDepositStake(f func(ctx context.Context, stakedAmount *big.Int) (common.Hash, error)) Option {
    66  	return func(mock *stakingContractMock) {
    67  		mock.depositStake = f
    68  	}
    69  }
    70  
    71  func WithGetStake(f func(ctx context.Context) (*big.Int, error)) Option {
    72  	return func(mock *stakingContractMock) {
    73  		mock.getStake = f
    74  	}
    75  }
    76  
    77  func WithWithdrawStake(f func(ctx context.Context) (common.Hash, error)) Option {
    78  	return func(mock *stakingContractMock) {
    79  		mock.withdrawAllStake = f
    80  	}
    81  }
    82  
    83  func WithMigrateStake(f func(ctx context.Context) (common.Hash, error)) Option {
    84  	return func(mock *stakingContractMock) {
    85  		mock.migrateStake = f
    86  	}
    87  }
    88  
    89  func WithIsFrozen(f func(ctx context.Context, block uint64) (bool, error)) Option {
    90  	return func(mock *stakingContractMock) {
    91  		mock.isFrozen = f
    92  	}
    93  }