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

     1  package calcium
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	enginemocks "github.com/projecteru2/core/engine/mocks"
     9  	lockmocks "github.com/projecteru2/core/lock/mocks"
    10  	resourcemocks "github.com/projecteru2/core/resource/mocks"
    11  	plugintypes "github.com/projecteru2/core/resource/plugins/types"
    12  	storemocks "github.com/projecteru2/core/store/mocks"
    13  	"github.com/projecteru2/core/types"
    14  
    15  	"github.com/cockroachdb/errors"
    16  	resourcetypes "github.com/projecteru2/core/resource/types"
    17  	"github.com/stretchr/testify/assert"
    18  	"github.com/stretchr/testify/mock"
    19  )
    20  
    21  func TestRemoveWorkload(t *testing.T) {
    22  	c := NewTestCluster()
    23  	ctx := context.Background()
    24  	lock := &lockmocks.DistributedLock{}
    25  	lock.On("Lock", mock.Anything).Return(ctx, nil)
    26  	lock.On("Unlock", mock.Anything).Return(nil)
    27  	store := c.store.(*storemocks.Store)
    28  	rmgr := c.rmgr.(*resourcemocks.Manager)
    29  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, nil, nil)
    30  	rmgr.On("SetNodeResourceUsage", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
    31  		resourcetypes.Resources{},
    32  		resourcetypes.Resources{},
    33  		nil,
    34  	)
    35  	rmgr.On("GetNodeMetrics", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]*plugintypes.Metrics{}, nil)
    36  
    37  	// failed by GetWorkload
    38  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    39  	ch, err := c.RemoveWorkload(ctx, []string{"xx"}, false)
    40  	assert.True(t, errors.Is(err, types.ErrMockError))
    41  	store.AssertExpectations(t)
    42  
    43  	// failed by GetNode
    44  	workload := &types.Workload{
    45  		ID:       "xx",
    46  		Name:     "test",
    47  		Nodename: "test",
    48  	}
    49  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
    50  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    51  	ch, err = c.RemoveWorkload(ctx, []string{"xx"}, false)
    52  	assert.NoError(t, err)
    53  	for r := range ch {
    54  		assert.False(t, r.Success)
    55  	}
    56  	time.Sleep(time.Second)
    57  	store.AssertExpectations(t)
    58  
    59  	// failed by Remove
    60  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
    61  	node := &types.Node{
    62  		NodeMeta: types.NodeMeta{
    63  			Name: "test",
    64  		},
    65  	}
    66  	store.On("GetNode", mock.Anything, mock.Anything).Return(node, nil)
    67  	store.On("RemoveWorkload", mock.Anything, mock.Anything).Return(types.ErrMockError).Twice()
    68  	store.On("ListNodeWorkloads", mock.Anything, mock.Anything, mock.Anything).Return(nil, types.ErrMockError)
    69  	ch, err = c.RemoveWorkload(ctx, []string{"xx"}, false)
    70  	assert.NoError(t, err)
    71  	for r := range ch {
    72  		assert.False(t, r.Success)
    73  	}
    74  	assert.NoError(t, c.doRemoveWorkloadSync(ctx, []string{"xx"}))
    75  	time.Sleep(time.Second)
    76  	store.AssertExpectations(t)
    77  
    78  	// success
    79  	engine := &enginemocks.API{}
    80  	workload.Engine = engine
    81  	engine.On("VirtualizationRemove", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
    82  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
    83  	store.On("RemoveWorkload", mock.Anything, mock.Anything).Return(nil)
    84  	ch, err = c.RemoveWorkload(ctx, []string{"xx"}, false)
    85  	assert.NoError(t, err)
    86  	for r := range ch {
    87  		assert.True(t, r.Success)
    88  	}
    89  	store.AssertExpectations(t)
    90  }