fix UT issue

projecteru2 · Oct 16, 2024 · a03dc8d · a03dc8d
1 parent 4e56b2b
commit a03dc8d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/cluster/calcium/node_test.go b/cluster/calcium/node_test.go
@@ -90,12 +90,14 @@ func TestRemoveNode(t *testing.T) {
 
 	// fail by store.RemoveNode
 	store.On("ListNodeWorkloads", mock.Anything, mock.Anything, mock.Anything).Return([]*types.Workload{}, nil)
+	store.On("SetNodeStatus", mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
 	store.On("RemoveNode", mock.Anything, mock.Anything).Return(types.ErrMockError).Once()
 	assert.Error(t, c.RemoveNode(ctx, name))
 
 	// success
+	store.On("SetNodeStatus", mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
 	store.On("RemoveNode", mock.Anything, mock.Anything).Return(nil)
-	store.On("SetNodeStatus", mock.Anything, mock.Anything, int64(-1)).Return(nil)
+	store.On("SetNodeStatus", mock.Anything, mock.Anything, int64(-1)).Return(nil).Once()
 	rmgr := c.rmgr.(*resourcemocks.Manager)
 	rmgr.On("RemoveNode", mock.Anything, mock.Anything).Return(nil)
 	assert.NoError(t, c.RemoveNode(ctx, name))