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

     1  package calcium
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	enginemocks "github.com/projecteru2/core/engine/mocks"
     8  	"github.com/projecteru2/core/lock"
     9  	lockmocks "github.com/projecteru2/core/lock/mocks"
    10  	resourcemocks "github.com/projecteru2/core/resource/mocks"
    11  	storemocks "github.com/projecteru2/core/store/mocks"
    12  	"github.com/projecteru2/core/types"
    13  
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/mock"
    16  )
    17  
    18  func TestDoLock(t *testing.T) {
    19  	c := NewTestCluster()
    20  	ctx := context.Background()
    21  	store := c.store.(*storemocks.Store)
    22  	// create lock failed
    23  	store.On("CreateLock", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    24  	_, _, err := c.doLock(ctx, "somename", 1)
    25  	assert.Error(t, err)
    26  
    27  	lock := &lockmocks.DistributedLock{}
    28  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
    29  	// lock failed
    30  	lock.On("Lock", mock.Anything).Return(context.Background(), types.ErrMockError).Once()
    31  	lock.On("Unlock", mock.Anything).Return(nil).Once()
    32  	_, _, err = c.doLock(ctx, "somename", 1)
    33  	assert.Error(t, err)
    34  	// success
    35  	lock.On("Lock", mock.Anything).Return(ctx, nil)
    36  	_, _, err = c.doLock(ctx, "somename", 1)
    37  	assert.NoError(t, err)
    38  }
    39  
    40  func TestDoUnlockAll(t *testing.T) {
    41  	c := NewTestCluster()
    42  	locks := map[string]lock.DistributedLock{}
    43  	lock := &lockmocks.DistributedLock{}
    44  	locks["somename"] = lock
    45  
    46  	// failed
    47  	lock.On("Unlock", mock.Anything).Return(types.ErrMockError)
    48  	c.doUnlockAll(context.Background(), locks)
    49  }
    50  
    51  func TestWithWorkloadsLocked(t *testing.T) {
    52  	c := NewTestCluster()
    53  	ctx := context.Background()
    54  	store := c.store.(*storemocks.Store)
    55  
    56  	lock := &lockmocks.DistributedLock{}
    57  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
    58  	lock.On("Unlock", mock.Anything).Return(nil)
    59  	// failed to get lock
    60  	lock.On("Lock", mock.Anything).Return(context.Background(), types.ErrMockError).Once()
    61  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{{}}, nil).Once()
    62  	err := c.withWorkloadsLocked(ctx, false, []string{"c1", "c2"}, func(ctx context.Context, workloads map[string]*types.Workload) error { return nil })
    63  	assert.Error(t, err)
    64  	// success
    65  	lock.On("Lock", mock.Anything).Return(ctx, nil)
    66  	// failed by getworkload
    67  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    68  	err = c.withWorkloadsLocked(ctx, false, []string{"c1", "c2"}, func(ctx context.Context, workloads map[string]*types.Workload) error { return nil })
    69  	assert.Error(t, err)
    70  	engine := &enginemocks.API{}
    71  	workload := &types.Workload{
    72  		ID:     "c1",
    73  		Engine: engine,
    74  	}
    75  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
    76  	// success
    77  	err = c.withWorkloadsLocked(ctx, false, []string{"c1", "c1"}, func(ctx context.Context, workloads map[string]*types.Workload) error {
    78  		assert.Len(t, workloads, 1)
    79  		return nil
    80  	})
    81  	assert.NoError(t, err)
    82  }
    83  
    84  func TestWithWorkloadLocked(t *testing.T) {
    85  	c := NewTestCluster()
    86  	ctx := context.Background()
    87  	store := c.store.(*storemocks.Store)
    88  
    89  	lock := &lockmocks.DistributedLock{}
    90  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
    91  	lock.On("Unlock", mock.Anything).Return(nil)
    92  	// failed to get lock
    93  	lock.On("Lock", mock.Anything).Return(context.Background(), types.ErrMockError).Once()
    94  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{{}}, nil).Once()
    95  	err := c.withWorkloadLocked(ctx, "c1", false, func(ctx context.Context, workload *types.Workload) error { return nil })
    96  	assert.Error(t, err)
    97  	// success
    98  	lock.On("Lock", mock.Anything).Return(ctx, nil)
    99  	// failed by getworkload
   100  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   101  	err = c.withWorkloadLocked(ctx, "c1", false, func(ctx context.Context, workload *types.Workload) error { return nil })
   102  	assert.Error(t, err)
   103  	engine := &enginemocks.API{}
   104  	workload := &types.Workload{
   105  		ID:     "c1",
   106  		Engine: engine,
   107  	}
   108  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
   109  	// success
   110  	err = c.withWorkloadLocked(ctx, "c1", false, func(ctx context.Context, workload *types.Workload) error {
   111  		assert.Equal(t, workload.ID, "c1")
   112  		return nil
   113  	})
   114  	assert.NoError(t, err)
   115  }
   116  
   117  func TestWithNodesPodLocked(t *testing.T) {
   118  	c := NewTestCluster()
   119  	ctx := context.Background()
   120  	store := c.store.(*storemocks.Store)
   121  	rmgr := c.rmgr.(*resourcemocks.Manager)
   122  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, nil, nil)
   123  
   124  	node1 := &types.Node{
   125  		NodeMeta: types.NodeMeta{
   126  			Name: "test",
   127  			Labels: map[string]string{
   128  				"eru": "1",
   129  			},
   130  			Podname: "test",
   131  		},
   132  		Available: true,
   133  	}
   134  	// failed by list nodes
   135  	store.On("GetNodesByPod", mock.Anything, mock.Anything).Return([]*types.Node{}, types.ErrMockError).Once()
   136  	err := c.withNodesPodLocked(ctx, &types.NodeFilter{Podname: "test", All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   137  	assert.Error(t, err)
   138  	store.On("GetNodesByPod", mock.Anything, mock.Anything).Return([]*types.Node{}, nil).Once()
   139  	// failed by filter
   140  	var ns map[string]*types.Node
   141  	err = c.withNodesPodLocked(ctx, &types.NodeFilter{Podname: "test", Labels: map[string]string{"eru": "2"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error {
   142  		ns = nodes
   143  		return nil
   144  	})
   145  	assert.NoError(t, err)
   146  	assert.Empty(t, ns)
   147  	store.On("GetNodesByPod", mock.Anything, mock.Anything).Return([]*types.Node{}, nil)
   148  	// failed by getnode
   149  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   150  	err = c.withNodesPodLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   151  	assert.Error(t, err)
   152  	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil).Once()
   153  	// failed by lock
   154  	lock := &lockmocks.DistributedLock{}
   155  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
   156  	lock.On("Unlock", mock.Anything).Return(nil)
   157  	// failed to get lock
   158  	lock.On("Lock", mock.Anything).Return(context.Background(), types.ErrMockError).Once()
   159  	err = c.withNodesPodLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   160  	assert.Error(t, err)
   161  	lock.On("Lock", mock.Anything).Return(ctx, nil)
   162  	// failed by get locked node
   163  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   164  	err = c.withNodesPodLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   165  	assert.Error(t, err)
   166  	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil)
   167  	// success
   168  	err = c.withNodesPodLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error {
   169  		assert.Len(t, nodes, 1)
   170  		return nil
   171  	})
   172  	assert.NoError(t, err)
   173  }
   174  
   175  func TestWithNodePodLocked(t *testing.T) {
   176  	c := NewTestCluster()
   177  	ctx := context.Background()
   178  	store := c.store.(*storemocks.Store)
   179  	rmgr := c.rmgr.(*resourcemocks.Manager)
   180  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, nil, nil)
   181  
   182  	node1 := &types.Node{
   183  		NodeMeta: types.NodeMeta{
   184  			Name: "test",
   185  			Labels: map[string]string{
   186  				"eru": "1",
   187  			},
   188  			Podname: "test",
   189  		},
   190  		Available: true,
   191  	}
   192  	// failed by lock
   193  	lock := &lockmocks.DistributedLock{}
   194  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
   195  	lock.On("Unlock", mock.Anything).Return(nil)
   196  	lock.On("Lock", mock.Anything).Return(ctx, nil)
   197  	// failed by get locked node
   198  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   199  	err := c.withNodePodLocked(ctx, "test", func(ctx context.Context, node *types.Node) error { return nil })
   200  	assert.Error(t, err)
   201  	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil)
   202  	// success
   203  	err = c.withNodePodLocked(ctx, "test", func(ctx context.Context, node *types.Node) error {
   204  		assert.Equal(t, node.Name, node1.Name)
   205  		return nil
   206  	})
   207  	assert.NoError(t, err)
   208  }
   209  
   210  func TestWithNodesOperationLocked(t *testing.T) {
   211  	c := NewTestCluster()
   212  	ctx := context.Background()
   213  	store := c.store.(*storemocks.Store)
   214  	rmgr := c.rmgr.(*resourcemocks.Manager)
   215  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, nil, nil)
   216  
   217  	node1 := &types.Node{
   218  		NodeMeta: types.NodeMeta{
   219  			Name: "test",
   220  			Labels: map[string]string{
   221  				"eru": "1",
   222  			},
   223  		},
   224  		Available: true,
   225  	}
   226  	// failed by list nodes
   227  	store.On("GetNodesByPod", mock.Anything, mock.Anything).Return([]*types.Node{}, types.ErrMockError).Once()
   228  	err := c.withNodesOperationLocked(ctx, &types.NodeFilter{Podname: "test", All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   229  	assert.Error(t, err)
   230  	store.On("GetNodesByPod", mock.Anything, mock.Anything).Return([]*types.Node{}, nil).Once()
   231  	// failed by filter
   232  	var ns map[string]*types.Node
   233  	err = c.withNodesOperationLocked(ctx, &types.NodeFilter{Podname: "test", Labels: map[string]string{"eru": "2"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error {
   234  		ns = nodes
   235  		return nil
   236  	})
   237  	assert.NoError(t, err)
   238  	assert.Empty(t, ns)
   239  	store.On("GetNodesByPod", mock.Anything, mock.Anything).Return([]*types.Node{}, nil)
   240  	// failed by getnode
   241  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   242  	err = c.withNodesOperationLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   243  	assert.Error(t, err)
   244  	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil).Once()
   245  	// failed by lock
   246  	lock := &lockmocks.DistributedLock{}
   247  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
   248  	lock.On("Unlock", mock.Anything).Return(nil)
   249  	// failed to get lock
   250  	lock.On("Lock", mock.Anything).Return(context.Background(), types.ErrMockError).Once()
   251  	err = c.withNodesOperationLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   252  	assert.Error(t, err)
   253  	lock.On("Lock", mock.Anything).Return(ctx, nil)
   254  	// failed by get locked node
   255  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   256  	err = c.withNodesOperationLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
   257  	assert.Error(t, err)
   258  	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil)
   259  	// success
   260  	err = c.withNodesOperationLocked(ctx, &types.NodeFilter{Podname: "test", Includes: []string{"test"}, All: false}, func(ctx context.Context, nodes map[string]*types.Node) error {
   261  		assert.Len(t, nodes, 1)
   262  		return nil
   263  	})
   264  	assert.NoError(t, err)
   265  }
   266  
   267  func TestWithNodeOperationLocked(t *testing.T) {
   268  	c := NewTestCluster()
   269  	ctx := context.Background()
   270  	store := c.store.(*storemocks.Store)
   271  	rmgr := c.rmgr.(*resourcemocks.Manager)
   272  	rmgr.On("GetNodeResourceInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, nil, nil)
   273  
   274  	node1 := &types.Node{
   275  		NodeMeta: types.NodeMeta{
   276  			Name: "test",
   277  			Labels: map[string]string{
   278  				"eru": "1",
   279  			},
   280  		},
   281  		Available: true,
   282  	}
   283  	// failed by lock
   284  	lock := &lockmocks.DistributedLock{}
   285  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
   286  	lock.On("Unlock", mock.Anything).Return(nil)
   287  	lock.On("Lock", mock.Anything).Return(ctx, nil)
   288  	// failed by get locked node
   289  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   290  	err := c.withNodeOperationLocked(ctx, "test", func(ctx context.Context, node *types.Node) error { return nil })
   291  	assert.Error(t, err)
   292  	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil)
   293  	// success
   294  	err = c.withNodeOperationLocked(ctx, "test", func(ctx context.Context, node *types.Node) error {
   295  		assert.Equal(t, node.Name, node1.Name)
   296  		return nil
   297  	})
   298  	assert.NoError(t, err)
   299  
   300  	err = c.withNodePodLocked(ctx, "test", func(ctx context.Context, node *types.Node) error {
   301  		return c.withNodeOperationLocked(ctx, node.Name, func(ctx context.Context, node *types.Node) error {
   302  			assert.Equal(t, node.Name, node1.Name)
   303  			return nil
   304  		})
   305  	})
   306  	assert.NoError(t, err)
   307  }