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

     1  package calcium
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/cockroachdb/errors"
    10  	enginemocks "github.com/projecteru2/core/engine/mocks"
    11  	enginetypes "github.com/projecteru2/core/engine/types"
    12  	lockmocks "github.com/projecteru2/core/lock/mocks"
    13  	resourcemocks "github.com/projecteru2/core/resource/mocks"
    14  	plugintypes "github.com/projecteru2/core/resource/plugins/types"
    15  	resourcetypes "github.com/projecteru2/core/resource/types"
    16  	storemocks "github.com/projecteru2/core/store/mocks"
    17  	"github.com/projecteru2/core/types"
    18  
    19  	"github.com/stretchr/testify/mock"
    20  	"github.com/stretchr/testify/require"
    21  )
    22  
    23  func TestHandleCreateWorkloadNoHandle(t *testing.T) {
    24  	c := NewTestCluster()
    25  	wal, err := enableWAL(c.config, c, c.store)
    26  	require.NoError(t, err)
    27  	c.wal = wal
    28  	rmgr := c.rmgr.(*resourcemocks.Manager)
    29  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
    30  		resourcetypes.Resources{},
    31  		resourcetypes.Resources{},
    32  		[]string{},
    33  		nil,
    34  	)
    35  
    36  	wrkid := "workload-id"
    37  	_, err = c.wal.Log(eventWorkloadCreated, &types.Workload{ID: wrkid, Nodename: "nodename"})
    38  	require.NoError(t, err)
    39  
    40  	wrk := &types.Workload{
    41  		ID: "wrkid",
    42  	}
    43  
    44  	store := c.store.(*storemocks.Store)
    45  	store.On("GetWorkload", mock.Anything, wrkid).Return(wrk, nil).Once()
    46  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, nil)
    47  
    48  	c.wal.Recover(context.Background())
    49  	store.AssertExpectations(t)
    50  
    51  	// Recovers nothing.
    52  	c.wal.Recover(context.Background())
    53  }
    54  
    55  func TestHandleCreateWorkloadError(t *testing.T) {
    56  	c := NewTestCluster()
    57  	wal, err := enableWAL(c.config, c, c.store)
    58  	require.NoError(t, err)
    59  	c.wal = wal
    60  	rmgr := c.rmgr.(*resourcemocks.Manager)
    61  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
    62  		resourcetypes.Resources{},
    63  		resourcetypes.Resources{},
    64  		[]string{},
    65  		nil,
    66  	)
    67  
    68  	engine := &enginemocks.API{}
    69  	node := &types.Node{
    70  		NodeMeta: types.NodeMeta{Name: "nodename"},
    71  		Engine:   engine,
    72  	}
    73  	wrkid := "workload-id"
    74  	_, err = c.wal.Log(eventWorkloadCreated, &types.Workload{ID: wrkid, Nodename: node.Name})
    75  	require.NoError(t, err)
    76  
    77  	wrk := &types.Workload{
    78  		ID:       wrkid,
    79  		Nodename: node.Name,
    80  	}
    81  
    82  	store := c.store.(*storemocks.Store)
    83  
    84  	err = errors.Wrapf(types.ErrInvaildCount, "keys: [%s]", wrkid)
    85  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(wrk, err).Once()
    86  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, err).Once()
    87  	c.wal.Recover(context.Background())
    88  	store.AssertExpectations(t)
    89  	engine.AssertExpectations(t)
    90  
    91  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(wrk, err).Once()
    92  	store.On("GetNode", mock.Anything, wrk.Nodename).Return(node, nil).Once()
    93  	engine.On("VirtualizationRemove", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(err).Once()
    94  	c.wal.Recover(context.Background())
    95  	store.AssertExpectations(t)
    96  	engine.AssertExpectations(t)
    97  
    98  	store.On("GetWorkload", mock.Anything, wrkid).Return(wrk, fmt.Errorf("err")).Once()
    99  	store.On("GetNode", mock.Anything, mock.Anything).Return(node, nil).Once()
   100  	engine.On("VirtualizationRemove", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(types.ErrWorkloadNotExists).Once()
   101  	c.wal.Recover(context.Background())
   102  	store.AssertExpectations(t)
   103  	engine.AssertExpectations(t)
   104  
   105  	// Nothing recovered.
   106  	c.wal.Recover(context.Background())
   107  }
   108  
   109  func TestHandleCreateWorkloadHandled(t *testing.T) {
   110  	c := NewTestCluster()
   111  	wal, err := enableWAL(c.config, c, c.store)
   112  	require.NoError(t, err)
   113  	c.wal = wal
   114  	rmgr := c.rmgr.(*resourcemocks.Manager)
   115  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
   116  		resourcetypes.Resources{},
   117  		resourcetypes.Resources{},
   118  		[]string{},
   119  		nil,
   120  	)
   121  
   122  	node := &types.Node{
   123  		NodeMeta: types.NodeMeta{Name: "nodename"},
   124  		Engine:   &enginemocks.API{},
   125  	}
   126  
   127  	wrkid := "workload-id"
   128  	_, err = c.wal.Log(eventWorkloadCreated, &types.Workload{ID: wrkid, Nodename: node.Name})
   129  	require.NoError(t, err)
   130  
   131  	wrk := &types.Workload{
   132  		ID:       wrkid,
   133  		Nodename: node.Name,
   134  		Engine:   nil, // explicitly set a nil as it will be evaluated by CreateWorkloadHandler.Handle.
   135  	}
   136  
   137  	store := c.store.(*storemocks.Store)
   138  
   139  	err = errors.Wrapf(types.ErrInvaildCount, "keys: [%s]", wrkid)
   140  	store.On("GetWorkload", mock.Anything, wrkid).Return(nil, err).Once()
   141  	store.On("GetNode", mock.Anything, wrk.Nodename).Return(node, nil)
   142  
   143  	eng, ok := node.Engine.(*enginemocks.API)
   144  	require.True(t, ok)
   145  	eng.On("VirtualizationRemove", mock.Anything, wrk.ID, true, true).
   146  		Return(nil).
   147  		Once()
   148  
   149  	c.wal.Recover(context.Background())
   150  	store.AssertExpectations(t)
   151  	eng.AssertExpectations(t)
   152  
   153  	// Recovers nothing.
   154  	c.wal.Recover(context.Background())
   155  }
   156  
   157  func TestHandleCreateLambda(t *testing.T) {
   158  	c := NewTestCluster()
   159  	wal, err := enableWAL(c.config, c, c.store)
   160  	require.NoError(t, err)
   161  	c.wal = wal
   162  	rmgr := c.rmgr.(*resourcemocks.Manager)
   163  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
   164  		resourcetypes.Resources{},
   165  		resourcetypes.Resources{},
   166  		[]string{},
   167  		nil,
   168  	)
   169  	rmgr.On("SetNodeResourceUsage", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
   170  		resourcetypes.Resources{},
   171  		resourcetypes.Resources{},
   172  		nil,
   173  	)
   174  	rmgr.On("GetNodeMetrics", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]*plugintypes.Metrics{}, nil)
   175  	rmgr.On("Remap", mock.Anything, mock.Anything, mock.Anything).Return(
   176  		resourcetypes.Resources{},
   177  		nil,
   178  	)
   179  
   180  	_, err = c.wal.Log(eventCreateLambda, "workloadid")
   181  	require.NoError(t, err)
   182  
   183  	node := &types.Node{
   184  		NodeMeta: types.NodeMeta{Name: "nodename"},
   185  		Engine:   &enginemocks.API{},
   186  	}
   187  	wrk := &types.Workload{
   188  		ID:       "workloadid",
   189  		Nodename: node.Name,
   190  		Engine:   node.Engine,
   191  	}
   192  
   193  	store := c.store.(*storemocks.Store)
   194  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   195  	c.wal.Recover(context.Background())
   196  	time.Sleep(500 * time.Millisecond)
   197  	store.AssertExpectations(t)
   198  
   199  	_, err = c.wal.Log(eventCreateLambda, "workloadid")
   200  	require.NoError(t, err)
   201  	store.On("GetWorkload", mock.Anything, mock.Anything).
   202  		Return(wrk, nil).
   203  		Once()
   204  	store.On("GetNode", mock.Anything, wrk.Nodename).
   205  		Return(node, nil)
   206  	eng := wrk.Engine.(*enginemocks.API)
   207  	eng.On("VirtualizationWait", mock.Anything, wrk.ID, "").Return(&enginetypes.VirtualizationWaitResult{Code: 0}, nil).Once()
   208  	eng.On("VirtualizationRemove", mock.Anything, wrk.ID, true, true).
   209  		Return(nil).
   210  		Once()
   211  	store.On("GetWorkloads", mock.Anything, []string{wrk.ID}).
   212  		Return([]*types.Workload{wrk}, nil).
   213  		Twice()
   214  	store.On("RemoveWorkload", mock.Anything, wrk).
   215  		Return(nil).
   216  		Once()
   217  	store.On("ListNodeWorkloads", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil).Once()
   218  	lock := &lockmocks.DistributedLock{}
   219  	lock.On("Lock", mock.Anything).Return(context.Background(), nil)
   220  	lock.On("Unlock", mock.Anything).Return(nil)
   221  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
   222  
   223  	c.wal.Recover(context.Background())
   224  	// Recovered nothing.
   225  	c.wal.Recover(context.Background())
   226  	time.Sleep(500 * time.Millisecond)
   227  	store.AssertExpectations(t)
   228  	eng.AssertExpectations(t)
   229  }