diff --git a/cluster/calcium/create_container_test.go b/cluster/calcium/create_container_test.go index 287742654..1d3e503d2 100644 --- a/cluster/calcium/create_container_test.go +++ b/cluster/calcium/create_container_test.go @@ -5,6 +5,7 @@ import ( "fmt" "testing" + "github.com/projecteru2/core/scheduler" "github.com/projecteru2/core/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" @@ -26,10 +27,16 @@ func TestPullImage(t *testing.T) { func TestCreateContainerWithMemPrior(t *testing.T) { initMockConfig() + ctx := context.Background() + mockStore.On("SaveProcessing", ctx, opts, mock.AnythingOfType("types.NodeInfo")).Return(nil) + mockStore.On("UpdateProcessing", ctx, opts, mock.AnythingOfType("string"), mock.AnythingOfType("int")).Return(nil) + mockStore.On("DeleteProcessing", ctx, opts, mock.AnythingOfType("types.NodeInfo")).Return(nil) + // Create Container with memory prior testlogF("Create containers with memory prior") - createCh, err := mockc.createContainerWithMemoryPrior(ctx, opts) + pod := &types.Pod{Favor: scheduler.MEMORY_PRIOR} + createCh, err := mockc.createContainer(ctx, opts, pod) assert.NoError(t, err) IDs := []string{} for msg := range createCh { @@ -74,14 +81,20 @@ func TestClean(t *testing.T) { func TestCreateContainerWithCPUPrior(t *testing.T) { initMockConfig() + ctx := context.Background() // update node mockStore.On("UpdateNode", mock.MatchedBy(func(input *types.Node) bool { return true })).Return(nil) - // Create Container with memory prior - testlogF("Create containers with memory prior") - createCh, err := mockc.createContainerWithCPUPrior(context.Background(), opts) + mockStore.On("SaveProcessing", ctx, opts, mock.AnythingOfType("types.NodeInfo")).Return(nil) + mockStore.On("UpdateProcessing", ctx, opts, mock.AnythingOfType("string"), mock.AnythingOfType("int")).Return(nil) + mockStore.On("DeleteProcessing", ctx, opts, mock.AnythingOfType("types.NodeInfo")).Return(nil) + + // Create Container with cpu prior + testlogF("Create containers with cpu prior") + pod := &types.Pod{Favor: scheduler.CPU_PRIOR} + createCh, err := mockc.createContainer(ctx, opts, pod) assert.NoError(t, err) IDs := []string{} for msg := range createCh { diff --git a/store/mock/store.go b/store/mock/store.go index 43b319a9c..5ccaa88a6 100644 --- a/store/mock/store.go +++ b/store/mock/store.go @@ -212,3 +212,21 @@ func (m *MockStore) ContainerDeployed(ctx context.Context, ID, appname, entrypoi args := m.Called(ID, appname, entrypoint, nodename, data) return args.Error(0) } + +// SaveProcessing save processing status in etcd +func (m *MockStore) SaveProcessing(ctx context.Context, opts *types.DeployOptions, nodeInfo types.NodeInfo) error { + args := m.Called(ctx, opts, nodeInfo) + return args.Error(0) +} + +// UpdateProcessing update processing status in etcd +func (m *MockStore) UpdateProcessing(ctx context.Context, opts *types.DeployOptions, nodename string, count int) error { + args := m.Called(ctx, opts, nodename, count) + return args.Error(0) +} + +// DeleteProcessing delete processing status in etcd +func (m *MockStore) DeleteProcessing(ctx context.Context, opts *types.DeployOptions, nodeInfo types.NodeInfo) error { + args := m.Called(ctx, opts, nodeInfo) + return args.Error(0) +}