github.com/ava-labs/avalanchego@v1.11.11/vms/rpcchainvm/with_context_vm_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package rpcchainvm
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/require"
    12  	"go.uber.org/mock/gomock"
    13  
    14  	"github.com/ava-labs/avalanchego/database/memdb"
    15  	"github.com/ava-labs/avalanchego/ids"
    16  	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
    17  	"github.com/ava-labs/avalanchego/snow/consensus/snowman/snowmanmock"
    18  	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
    19  	"github.com/ava-labs/avalanchego/snow/engine/snowman/block/blockmock"
    20  	"github.com/ava-labs/avalanchego/snow/snowtest"
    21  )
    22  
    23  var (
    24  	_ block.ChainVM                      = ContextEnabledVMMock{}
    25  	_ block.BuildBlockWithContextChainVM = ContextEnabledVMMock{}
    26  
    27  	_ snowman.Block           = ContextEnabledBlockMock{}
    28  	_ block.WithVerifyContext = ContextEnabledBlockMock{}
    29  
    30  	blockContext = &block.Context{
    31  		PChainHeight: 1,
    32  	}
    33  
    34  	blkID    = ids.ID{1}
    35  	parentID = ids.ID{0}
    36  	blkBytes = []byte{0}
    37  )
    38  
    39  type ContextEnabledVMMock struct {
    40  	*blockmock.ChainVM
    41  	*blockmock.BuildBlockWithContextChainVM
    42  }
    43  
    44  type ContextEnabledBlockMock struct {
    45  	*snowmanmock.Block
    46  	*blockmock.WithVerifyContext
    47  }
    48  
    49  func contextEnabledTestPlugin(t *testing.T, loadExpectations bool) block.ChainVM {
    50  	// test key is "contextTestKey"
    51  
    52  	// create mock
    53  	ctrl := gomock.NewController(t)
    54  	ctxVM := ContextEnabledVMMock{
    55  		ChainVM:                      blockmock.NewChainVM(ctrl),
    56  		BuildBlockWithContextChainVM: blockmock.NewBuildBlockWithContextChainVM(ctrl),
    57  	}
    58  
    59  	if loadExpectations {
    60  		ctxBlock := ContextEnabledBlockMock{
    61  			Block:             snowmanmock.NewBlock(ctrl),
    62  			WithVerifyContext: blockmock.NewWithVerifyContext(ctrl),
    63  		}
    64  		gomock.InOrder(
    65  			// Initialize
    66  			ctxVM.ChainVM.EXPECT().Initialize(
    67  				gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(),
    68  				gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(),
    69  				gomock.Any(),
    70  			).Return(nil).Times(1),
    71  			ctxVM.ChainVM.EXPECT().LastAccepted(gomock.Any()).Return(preSummaryBlk.ID(), nil).Times(1),
    72  			ctxVM.ChainVM.EXPECT().GetBlock(gomock.Any(), gomock.Any()).Return(preSummaryBlk, nil).Times(1),
    73  
    74  			// BuildBlockWithContext
    75  			ctxVM.BuildBlockWithContextChainVM.EXPECT().BuildBlockWithContext(gomock.Any(), blockContext).Return(ctxBlock, nil).Times(1),
    76  			ctxBlock.WithVerifyContext.EXPECT().ShouldVerifyWithContext(gomock.Any()).Return(true, nil).Times(1),
    77  			ctxBlock.Block.EXPECT().ID().Return(blkID).Times(1),
    78  			ctxBlock.Block.EXPECT().Parent().Return(parentID).Times(1),
    79  			ctxBlock.Block.EXPECT().Bytes().Return(blkBytes).Times(1),
    80  			ctxBlock.Block.EXPECT().Height().Return(uint64(1)).Times(1),
    81  			ctxBlock.Block.EXPECT().Timestamp().Return(time.Now()).Times(1),
    82  
    83  			// VerifyWithContext
    84  			ctxVM.ChainVM.EXPECT().ParseBlock(gomock.Any(), blkBytes).Return(ctxBlock, nil).Times(1),
    85  			ctxBlock.WithVerifyContext.EXPECT().VerifyWithContext(gomock.Any(), blockContext).Return(nil).Times(1),
    86  			ctxBlock.Block.EXPECT().Timestamp().Return(time.Now()).Times(1),
    87  		)
    88  	}
    89  
    90  	return ctxVM
    91  }
    92  
    93  func TestContextVMSummary(t *testing.T) {
    94  	require := require.New(t)
    95  	testKey := contextTestKey
    96  
    97  	// Create and start the plugin
    98  	vm := buildClientHelper(require, testKey)
    99  	defer vm.runtime.Stop(context.Background())
   100  
   101  	ctx := snowtest.Context(t, snowtest.CChainID)
   102  
   103  	require.NoError(vm.Initialize(context.Background(), ctx, memdb.New(), nil, nil, nil, nil, nil, nil))
   104  
   105  	blkIntf, err := vm.BuildBlockWithContext(context.Background(), blockContext)
   106  	require.NoError(err)
   107  
   108  	blk, ok := blkIntf.(block.WithVerifyContext)
   109  	require.True(ok)
   110  
   111  	shouldVerify, err := blk.ShouldVerifyWithContext(context.Background())
   112  	require.NoError(err)
   113  	require.True(shouldVerify)
   114  
   115  	require.NoError(blk.VerifyWithContext(context.Background(), blockContext))
   116  }