github.com/ava-labs/avalanchego@v1.11.11/vms/registry/vm_registry_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 registry
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  	"go.uber.org/mock/gomock"
    12  
    13  	"github.com/ava-labs/avalanchego/ids"
    14  	"github.com/ava-labs/avalanchego/vms"
    15  	"github.com/ava-labs/avalanchego/vms/registry/registrymock"
    16  	"github.com/ava-labs/avalanchego/vms/vmsmock"
    17  )
    18  
    19  var (
    20  	id1 = ids.GenerateTestID()
    21  	id2 = ids.GenerateTestID()
    22  	id3 = ids.GenerateTestID()
    23  	id4 = ids.GenerateTestID()
    24  )
    25  
    26  // Tests the happy case where Reload succeeds.
    27  func TestReload_Success(t *testing.T) {
    28  	require := require.New(t)
    29  
    30  	resources := initVMRegistryTest(t)
    31  
    32  	factory1 := vmsmock.NewFactory(resources.ctrl)
    33  	factory2 := vmsmock.NewFactory(resources.ctrl)
    34  	factory3 := vmsmock.NewFactory(resources.ctrl)
    35  	factory4 := vmsmock.NewFactory(resources.ctrl)
    36  
    37  	registeredVms := map[ids.ID]vms.Factory{
    38  		id1: factory1,
    39  		id2: factory2,
    40  	}
    41  
    42  	unregisteredVms := map[ids.ID]vms.Factory{
    43  		id3: factory3,
    44  		id4: factory4,
    45  	}
    46  
    47  	resources.mockVMGetter.EXPECT().
    48  		Get().
    49  		Times(1).
    50  		Return(registeredVms, unregisteredVms, nil)
    51  	resources.mockVMManager.EXPECT().
    52  		RegisterFactory(gomock.Any(), id3, factory3).
    53  		Times(1).
    54  		Return(nil)
    55  	resources.mockVMManager.EXPECT().
    56  		RegisterFactory(gomock.Any(), id4, factory4).
    57  		Times(1).
    58  		Return(nil)
    59  
    60  	installedVMs, failedVMs, err := resources.vmRegistry.Reload(context.Background())
    61  	require.NoError(err)
    62  	require.ElementsMatch([]ids.ID{id3, id4}, installedVMs)
    63  	require.Empty(failedVMs)
    64  }
    65  
    66  // Tests that we fail if we're not able to get the vms on disk
    67  func TestReload_GetNewVMsFails(t *testing.T) {
    68  	require := require.New(t)
    69  
    70  	resources := initVMRegistryTest(t)
    71  
    72  	resources.mockVMGetter.EXPECT().Get().Times(1).Return(nil, nil, errTest)
    73  
    74  	installedVMs, failedVMs, err := resources.vmRegistry.Reload(context.Background())
    75  	require.ErrorIs(err, errTest)
    76  	require.Empty(installedVMs)
    77  	require.Empty(failedVMs)
    78  }
    79  
    80  // Tests that if we fail to register a VM, we fail.
    81  func TestReload_PartialRegisterFailure(t *testing.T) {
    82  	require := require.New(t)
    83  
    84  	resources := initVMRegistryTest(t)
    85  
    86  	factory1 := vmsmock.NewFactory(resources.ctrl)
    87  	factory2 := vmsmock.NewFactory(resources.ctrl)
    88  	factory3 := vmsmock.NewFactory(resources.ctrl)
    89  	factory4 := vmsmock.NewFactory(resources.ctrl)
    90  
    91  	registeredVms := map[ids.ID]vms.Factory{
    92  		id1: factory1,
    93  		id2: factory2,
    94  	}
    95  
    96  	unregisteredVms := map[ids.ID]vms.Factory{
    97  		id3: factory3,
    98  		id4: factory4,
    99  	}
   100  
   101  	resources.mockVMGetter.EXPECT().
   102  		Get().
   103  		Times(1).
   104  		Return(registeredVms, unregisteredVms, nil)
   105  	resources.mockVMManager.EXPECT().
   106  		RegisterFactory(gomock.Any(), id3, factory3).
   107  		Times(1).
   108  		Return(errTest)
   109  	resources.mockVMManager.EXPECT().
   110  		RegisterFactory(gomock.Any(), id4, factory4).
   111  		Times(1).
   112  		Return(nil)
   113  
   114  	installedVMs, failedVMs, err := resources.vmRegistry.Reload(context.Background())
   115  	require.NoError(err)
   116  	require.Len(failedVMs, 1)
   117  	require.ErrorIs(failedVMs[id3], errTest)
   118  	require.Len(installedVMs, 1)
   119  	require.Equal(id4, installedVMs[0])
   120  }
   121  
   122  type registryTestResources struct {
   123  	ctrl          *gomock.Controller
   124  	mockVMGetter  *registrymock.VMGetter
   125  	mockVMManager *vmsmock.Manager
   126  	vmRegistry    VMRegistry
   127  }
   128  
   129  func initVMRegistryTest(t *testing.T) *registryTestResources {
   130  	ctrl := gomock.NewController(t)
   131  
   132  	mockVMGetter := registrymock.NewVMGetter(ctrl)
   133  	mockVMManager := vmsmock.NewManager(ctrl)
   134  
   135  	vmRegistry := NewVMRegistry(
   136  		VMRegistryConfig{
   137  			VMGetter:  mockVMGetter,
   138  			VMManager: mockVMManager,
   139  		},
   140  	)
   141  
   142  	return &registryTestResources{
   143  		ctrl:          ctrl,
   144  		mockVMGetter:  mockVMGetter,
   145  		mockVMManager: mockVMManager,
   146  		vmRegistry:    vmRegistry,
   147  	}
   148  }