fix realloc tests

projecteru2 · Jul 8, 2022 · 766c9b7 · 766c9b7
1 parent fe0744a
commit 766c9b7
Showing 1 changed file with 15 additions and 15 deletions.
diff --git a/cluster/calcium/realloc_test.go b/cluster/calcium/realloc_test.go
@@ -7,6 +7,7 @@ import (
 	enginemocks "github.com/projecteru2/core/engine/mocks"
 	enginetypes "github.com/projecteru2/core/engine/types"
 	lockmocks "github.com/projecteru2/core/lock/mocks"
+	resourcemocks "github.com/projecteru2/core/resources/mocks"
 	storemocks "github.com/projecteru2/core/store/mocks"
 	"github.com/projecteru2/core/types"
 
@@ -18,8 +19,8 @@ func TestRealloc(t *testing.T) {
 	c := NewTestCluster()
 	ctx := context.Background()
 	store := c.store.(*storemocks.Store)
+	rmgr := c.rmgr.(*resourcemocks.Manager)
 	c.config.Scheduler.ShareBase = 100
-	//	plugin := c.resource.GetPlugins()[0].(*resourcemocks.Plugin)
 
 	lock := &lockmocks.DistributedLock{}
 	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
@@ -60,12 +61,7 @@ func TestRealloc(t *testing.T) {
 	assert.EqualError(t, err, "ETCD must be set")
 	store.AssertExpectations(t)
 	store.On("GetNode", mock.Anything, "node1").Return(node1, nil)
-	//	plugin.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything).Return(&resources.GetNodeResourceInfoResponse{
-	//		ResourceInfo: &resources.NodeResourceInfo{
-	//			Capacity: types.NodeResourceArgs{},
-	//			Usage:    types.NodeResourceArgs{},
-	//		},
-	//	}, nil)
+	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, nil, nil)
 
 	// failed by lock
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(nil, types.ErrNoETCD).Once()
@@ -76,14 +72,18 @@ func TestRealloc(t *testing.T) {
 	store.On("GetWorkloads", mock.Anything, []string{"c1"}).Return(newC1, nil)
 
 	// failed by plugin
-	//	plugin.On("GetReallocArgs", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, context.DeadlineExceeded).Once()
-	//	err = c.ReallocResource(ctx, opts)
-	//	assert.Error(t, err)
-	//	plugin.On("GetReallocArgs", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&resources.GetReallocArgsResponse{}, nil)
-	//	plugin.On("SetNodeResourceUsage", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&resources.SetNodeResourceUsageResponse{
-	//		Before: types.NodeResourceArgs{},
-	//		After:  types.NodeResourceArgs{},
-	//	}, nil)
+	rmgr.On("Realloc", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
+		types.EngineArgs{}, nil, nil, types.ErrNoETCD,
+	).Once()
+	err = c.ReallocResource(ctx, opts)
+	assert.Error(t, err)
+	rmgr.On("Realloc", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
+		types.EngineArgs{},
+		map[string]types.WorkloadResourceArgs{},
+		map[string]types.WorkloadResourceArgs{},
+		nil,
+	)
+	rmgr.On("RollbackRealloc", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 
 	// failed by UpdateWorkload
 	store.On("UpdateWorkload", mock.Anything, mock.Anything).Return(types.ErrNoETCD).Once()