github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/cluster/calcium/realloc_test.go (about)

     1  package calcium
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/cockroachdb/errors"
     8  	enginemocks "github.com/projecteru2/core/engine/mocks"
     9  	enginetypes "github.com/projecteru2/core/engine/types"
    10  	lockmocks "github.com/projecteru2/core/lock/mocks"
    11  	resourcemocks "github.com/projecteru2/core/resource/mocks"
    12  	plugintypes "github.com/projecteru2/core/resource/plugins/types"
    13  	resourcetypes "github.com/projecteru2/core/resource/types"
    14  	storemocks "github.com/projecteru2/core/store/mocks"
    15  	"github.com/projecteru2/core/types"
    16  
    17  	"github.com/stretchr/testify/assert"
    18  	"github.com/stretchr/testify/mock"
    19  )
    20  
    21  func TestRealloc(t *testing.T) {
    22  	c := NewTestCluster()
    23  	ctx := context.Background()
    24  	store := c.store.(*storemocks.Store)
    25  	rmgr := c.rmgr.(*resourcemocks.Manager)
    26  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, nil, nil)
    27  	rmgr.On("GetNodeMetrics", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]*plugintypes.Metrics{}, nil)
    28  	c.config.Scheduler.ShareBase = 100
    29  
    30  	lock := &lockmocks.DistributedLock{}
    31  	lock.On("Lock", mock.Anything).Return(ctx, nil)
    32  	lock.On("Unlock", mock.Anything).Return(nil)
    33  
    34  	engine := &enginemocks.API{}
    35  	engine.On("VirtualizationInspect", mock.Anything, mock.Anything).Return(&enginetypes.VirtualizationInfo{}, nil)
    36  
    37  	node1 := &types.Node{
    38  		NodeMeta: types.NodeMeta{
    39  			Name:     "node1",
    40  			Endpoint: "http://1.1.1.1:1",
    41  		},
    42  		Engine: engine,
    43  	}
    44  
    45  	newC1 := func(context.Context, []string) []*types.Workload {
    46  		return []*types.Workload{
    47  			{
    48  				ID:        "c1",
    49  				Podname:   "p1",
    50  				Engine:    engine,
    51  				Resources: resourcetypes.Resources{},
    52  				Nodename:  "node1",
    53  			},
    54  		}
    55  	}
    56  
    57  	store.On("GetWorkload", mock.Anything, "c1").Return(newC1(context.Background(), nil)[0], nil)
    58  	opts := &types.ReallocOptions{
    59  		ID:        "c1",
    60  		Resources: resourcetypes.Resources{},
    61  	}
    62  
    63  	// failed by GetNode
    64  	store.On("GetNode", mock.Anything, "node1").Return(nil, types.ErrMockError).Once()
    65  	err := c.ReallocResource(ctx, opts)
    66  	assert.True(t, errors.Is(err, types.ErrMockError))
    67  	store.AssertExpectations(t)
    68  	store.On("GetNode", mock.Anything, "node1").Return(node1, nil)
    69  
    70  	// failed by lock
    71  	store.On("CreateLock", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    72  	err = c.ReallocResource(ctx, opts)
    73  	assert.True(t, errors.Is(err, types.ErrMockError))
    74  	store.AssertExpectations(t)
    75  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
    76  	store.On("GetWorkloads", mock.Anything, []string{"c1"}).Return(newC1, nil)
    77  
    78  	// failed by plugin
    79  	rmgr.On("Realloc", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
    80  		resourcetypes.Resources{}, nil, nil, types.ErrMockError,
    81  	).Once()
    82  	err = c.ReallocResource(ctx, opts)
    83  	assert.Error(t, err)
    84  	rmgr.On("Realloc", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
    85  		resourcetypes.Resources{},
    86  		resourcetypes.Resources{},
    87  		resourcetypes.Resources{},
    88  		nil,
    89  	)
    90  	rmgr.On("RollbackRealloc", mock.Anything, mock.Anything, mock.Anything).Return(nil)
    91  
    92  	// failed by UpdateWorkload
    93  	store.On("UpdateWorkload", mock.Anything, mock.Anything).Return(types.ErrMockError).Once()
    94  	err = c.ReallocResource(ctx, opts)
    95  	assert.True(t, errors.Is(err, types.ErrMockError))
    96  	store.AssertExpectations(t)
    97  	store.On("UpdateWorkload", mock.Anything, mock.Anything).Return(nil)
    98  
    99  	// failed by virtualization update resource
   100  	engine.On("VirtualizationUpdateResource", mock.Anything, mock.Anything, mock.Anything).Return(types.ErrNilEngine).Once()
   101  	err = c.ReallocResource(ctx, opts)
   102  	assert.ErrorIs(t, err, types.ErrNilEngine)
   103  	engine.On("VirtualizationUpdateResource", mock.Anything, mock.Anything, mock.Anything).Return(nil)
   104  
   105  	// success
   106  	store.On("ListNodeWorkloads", mock.Anything, mock.Anything, mock.Anything).Return(nil, types.ErrMockError)
   107  	err = c.ReallocResource(ctx, opts)
   108  	assert.Nil(t, err)
   109  }