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

     1  package calcium
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	storemocks "github.com/projecteru2/core/store/mocks"
     8  	"github.com/projecteru2/core/types"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/mock"
    12  )
    13  
    14  func TestGetNodeStatus(t *testing.T) {
    15  	c := NewTestCluster()
    16  	ctx := context.Background()
    17  	store := c.store.(*storemocks.Store)
    18  
    19  	ns := &types.NodeStatus{
    20  		Nodename: "test",
    21  		Podname:  "test",
    22  		Alive:    true,
    23  	}
    24  	// failed
    25  	store.On("GetNodeStatus", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    26  	_, err := c.GetNodeStatus(ctx, "test")
    27  	assert.Error(t, err)
    28  
    29  	store.On("GetNodeStatus", mock.Anything, mock.Anything).Return(ns, nil)
    30  	s, err := c.GetNodeStatus(ctx, "test")
    31  	assert.NoError(t, err)
    32  	assert.Equal(t, s.Nodename, "test")
    33  	assert.True(t, s.Alive)
    34  	store.AssertExpectations(t)
    35  }
    36  
    37  func TestSetNodeStatus(t *testing.T) {
    38  	c := NewTestCluster()
    39  	ctx := context.Background()
    40  	store := c.store.(*storemocks.Store)
    41  
    42  	node := &types.Node{
    43  		NodeMeta: types.NodeMeta{
    44  			Name:     "testname",
    45  			Endpoint: "ep",
    46  		},
    47  	}
    48  	// failed by GetNode
    49  	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    50  	assert.Error(t, c.SetNodeStatus(ctx, node.Name, 10))
    51  	// failed by SetWorkloadStatus
    52  	store.On("GetNode", mock.Anything, mock.Anything).Return(node, nil)
    53  	store.On("SetNodeStatus",
    54  		mock.Anything,
    55  		mock.Anything,
    56  		mock.Anything,
    57  	).Return(types.ErrMockError).Once()
    58  	assert.Error(t, c.SetNodeStatus(ctx, node.Name, 10))
    59  	// success
    60  	store.On("SetNodeStatus",
    61  		mock.Anything,
    62  		mock.Anything,
    63  		mock.Anything,
    64  	).Return(nil)
    65  	assert.NoError(t, c.SetNodeStatus(ctx, node.Name, 10))
    66  	store.AssertExpectations(t)
    67  }
    68  
    69  func TestNodeStatusStream(t *testing.T) {
    70  	c := NewTestCluster()
    71  	ctx := context.Background()
    72  	dataCh := make(chan *types.NodeStatus)
    73  	store := c.store.(*storemocks.Store)
    74  	store.On("NodeStatusStream", mock.Anything).Return(dataCh)
    75  
    76  	go func() {
    77  		defer close(dataCh)
    78  		msg := &types.NodeStatus{
    79  			Alive: true,
    80  		}
    81  		dataCh <- msg
    82  	}()
    83  
    84  	for c := range c.NodeStatusStream(ctx) {
    85  		assert.True(t, c.Alive)
    86  	}
    87  	store.AssertExpectations(t)
    88  }
    89  
    90  func TestGetWorkloadsStatus(t *testing.T) {
    91  	c := NewTestCluster()
    92  	ctx := context.Background()
    93  	store := c.store.(*storemocks.Store)
    94  	cs := &types.StatusMeta{}
    95  
    96  	// failed GetWorkloadStatus
    97  	store.On("GetWorkloadStatus", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    98  	_, err := c.GetWorkloadsStatus(ctx, []string{"a"})
    99  	assert.Error(t, err)
   100  	// success
   101  	store.On("GetWorkloadStatus", mock.Anything, mock.Anything).Return(cs, nil)
   102  	r, err := c.GetWorkloadsStatus(ctx, []string{"a"})
   103  	assert.NoError(t, err)
   104  	assert.Len(t, r, 1)
   105  	store.AssertExpectations(t)
   106  }
   107  
   108  func TestSetWorkloadsStatus(t *testing.T) {
   109  	c := NewTestCluster()
   110  	ctx := context.Background()
   111  	store := c.store.(*storemocks.Store)
   112  
   113  	// no meta, generate by id
   114  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
   115  	_, err := c.SetWorkloadsStatus(ctx, []*types.StatusMeta{{ID: "123"}}, nil)
   116  	assert.Error(t, err)
   117  
   118  	// failed by workload name
   119  	workload := &types.Workload{
   120  		ID:   "123",
   121  		Name: "invalid",
   122  	}
   123  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(workload, nil).Once()
   124  	_, err = c.SetWorkloadsStatus(ctx, []*types.StatusMeta{{ID: "123"}}, nil)
   125  	assert.Error(t, err)
   126  
   127  	// failed by SetworkloadStatus
   128  	workload.Name = "a_b_c"
   129  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(workload, nil)
   130  	store.On("SetWorkloadStatus",
   131  		mock.Anything,
   132  		mock.Anything,
   133  		mock.Anything,
   134  	).Return(types.ErrMockError).Once()
   135  	_, err = c.SetWorkloadsStatus(ctx, []*types.StatusMeta{{ID: "123"}}, nil)
   136  	assert.Error(t, err)
   137  
   138  	// success
   139  	store.On("SetWorkloadStatus",
   140  		mock.Anything,
   141  		mock.Anything,
   142  		mock.Anything,
   143  	).Return(nil)
   144  	r, err := c.SetWorkloadsStatus(ctx, []*types.StatusMeta{{ID: "123"}}, nil)
   145  	assert.NoError(t, err)
   146  	assert.Len(t, r, 1)
   147  	store.AssertExpectations(t)
   148  }
   149  
   150  func TestWorkloadStatusStream(t *testing.T) {
   151  	c := NewTestCluster()
   152  	ctx := context.Background()
   153  	dataCh := make(chan *types.WorkloadStatus)
   154  	store := c.store.(*storemocks.Store)
   155  
   156  	store.On("WorkloadStatusStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(dataCh)
   157  	go func() {
   158  		defer close(dataCh)
   159  		msg := &types.WorkloadStatus{
   160  			Delete: true,
   161  		}
   162  		dataCh <- msg
   163  	}()
   164  
   165  	for c := range c.WorkloadStatusStream(ctx, "", "", "", nil) {
   166  		assert.Equal(t, c.Delete, true)
   167  	}
   168  }