code.vegaprotocol.io/vega@v0.79.0/core/referral/helpers_for_test.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 referral_test
    17  
    18  import (
    19  	"context"
    20  	"testing"
    21  	"time"
    22  
    23  	"code.vegaprotocol.io/vega/core/events"
    24  	"code.vegaprotocol.io/vega/core/integration/stubs"
    25  	"code.vegaprotocol.io/vega/core/referral"
    26  	"code.vegaprotocol.io/vega/core/referral/mocks"
    27  	"code.vegaprotocol.io/vega/core/snapshot"
    28  	"code.vegaprotocol.io/vega/core/stats"
    29  	"code.vegaprotocol.io/vega/core/types"
    30  	vgcrypto "code.vegaprotocol.io/vega/libs/crypto"
    31  	vgrand "code.vegaprotocol.io/vega/libs/rand"
    32  	"code.vegaprotocol.io/vega/logging"
    33  	"code.vegaprotocol.io/vega/paths"
    34  	vegapb "code.vegaprotocol.io/vega/protos/vega"
    35  
    36  	"github.com/golang/mock/gomock"
    37  	"github.com/stretchr/testify/assert"
    38  	"github.com/stretchr/testify/require"
    39  )
    40  
    41  type testEngine struct {
    42  	engine                *referral.SnapshottedEngine
    43  	broker                *mocks.MockBroker
    44  	timeSvc               *mocks.MockTimeService
    45  	marketActivityTracker *mocks.MockMarketActivityTracker
    46  	staking               *mocks.MockStakingBalances
    47  	currentEpoch          uint64
    48  }
    49  
    50  func newPartyID(t *testing.T) types.PartyID {
    51  	t.Helper()
    52  
    53  	return types.PartyID(vgrand.RandomStr(5))
    54  }
    55  
    56  func newSetID(t *testing.T) types.ReferralSetID {
    57  	t.Helper()
    58  
    59  	return types.ReferralSetID(vgcrypto.RandomHash())
    60  }
    61  
    62  func newSnapshotEngine(t *testing.T, vegaPath paths.Paths, now time.Time, engine *referral.SnapshottedEngine) *snapshot.Engine {
    63  	t.Helper()
    64  
    65  	log := logging.NewTestLogger()
    66  	timeService := stubs.NewTimeStub()
    67  	timeService.SetTime(now)
    68  	statsData := stats.New(log, stats.NewDefaultConfig())
    69  	config := snapshot.DefaultConfig()
    70  
    71  	snapshotEngine, err := snapshot.NewEngine(vegaPath, config, log, timeService, statsData.Blockchain)
    72  	require.NoError(t, err)
    73  
    74  	snapshotEngine.AddProviders(engine)
    75  
    76  	return snapshotEngine
    77  }
    78  
    79  func newEngine(t *testing.T) *testEngine {
    80  	t.Helper()
    81  
    82  	ctrl := gomock.NewController(t)
    83  
    84  	broker := mocks.NewMockBroker(ctrl)
    85  	timeSvc := mocks.NewMockTimeService(ctrl)
    86  	mat := mocks.NewMockMarketActivityTracker(ctrl)
    87  	staking := mocks.NewMockStakingBalances(ctrl)
    88  
    89  	engine := referral.NewSnapshottedEngine(broker, timeSvc, mat, staking)
    90  
    91  	engine.OnEpochRestore(context.Background(), types.Epoch{
    92  		Seq:    10,
    93  		Action: vegapb.EpochAction_EPOCH_ACTION_START,
    94  	})
    95  
    96  	return &testEngine{
    97  		engine:                engine,
    98  		broker:                broker,
    99  		timeSvc:               timeSvc,
   100  		marketActivityTracker: mat,
   101  		currentEpoch:          10,
   102  		staking:               staking,
   103  	}
   104  }
   105  
   106  func nextEpoch(t *testing.T, ctx context.Context, te *testEngine, startEpochTime time.Time) {
   107  	t.Helper()
   108  
   109  	te.engine.OnEpoch(ctx, types.Epoch{
   110  		Seq:     te.currentEpoch,
   111  		Action:  vegapb.EpochAction_EPOCH_ACTION_END,
   112  		EndTime: startEpochTime.Add(-1 * time.Second),
   113  	})
   114  
   115  	te.currentEpoch += 1
   116  
   117  	te.engine.OnEpoch(ctx, types.Epoch{
   118  		Seq:       te.currentEpoch,
   119  		Action:    vegapb.EpochAction_EPOCH_ACTION_START,
   120  		StartTime: startEpochTime,
   121  	})
   122  }
   123  
   124  func expectReferralProgramStartedEvent(t *testing.T, engine *testEngine) {
   125  	t.Helper()
   126  
   127  	engine.broker.EXPECT().Send(gomock.Any()).Do(func(evt events.Event) {
   128  		_, ok := evt.(*events.ReferralProgramStarted)
   129  		assert.True(t, ok, "Event should be a ReferralProgramStarted, but is %T", evt)
   130  	}).Times(1)
   131  }
   132  
   133  func expectReferralProgramEndedEvent(t *testing.T, engine *testEngine) *gomock.Call {
   134  	t.Helper()
   135  
   136  	return engine.broker.EXPECT().Send(gomock.Any()).Do(func(evt events.Event) {
   137  		_, ok := evt.(*events.ReferralProgramEnded)
   138  		assert.True(t, ok, "Event should be a ReferralProgramEnded, but is %T", evt)
   139  	}).Times(1)
   140  }
   141  
   142  func expectReferralProgramUpdatedEvent(t *testing.T, engine *testEngine) *gomock.Call {
   143  	t.Helper()
   144  
   145  	return engine.broker.EXPECT().Send(gomock.Any()).Do(func(evt events.Event) {
   146  		_, ok := evt.(*events.ReferralProgramUpdated)
   147  		assert.True(t, ok, "Event should be a ReferralProgramUpdated, but is %T", evt)
   148  	}).Times(1)
   149  }
   150  
   151  func expectReferralSetStatsUpdatedEvent(t *testing.T, engine *testEngine, times int) *gomock.Call {
   152  	t.Helper()
   153  
   154  	return engine.broker.EXPECT().Send(gomock.Any()).Do(func(evt events.Event) {
   155  		_, ok := evt.(*events.ReferralSetStatsUpdated)
   156  		assert.True(t, ok, "Event should be a ReferralSetStatsUpdated, but is %T", evt)
   157  	}).Times(times)
   158  }