github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/store/etcdv3/workload_test.go (about)

     1  package etcdv3
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/projecteru2/core/types"
    11  
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  func TestAddORUpdateWorkload(t *testing.T) {
    16  	m := NewMercury(t)
    17  	ctx := context.Background()
    18  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
    19  	name := "test_app_1"
    20  	nodename := "n1"
    21  	podname := "test"
    22  	workload := &types.Workload{
    23  		ID:       ID,
    24  		Nodename: nodename,
    25  		Podname:  podname,
    26  		Name:     "a",
    27  	}
    28  	// failed by name
    29  	err := m.AddWorkload(ctx, workload, nil)
    30  	assert.Error(t, err)
    31  	workload.Name = name
    32  	// fail update
    33  	err = m.UpdateWorkload(ctx, workload)
    34  	assert.Error(t, err)
    35  	// success create
    36  	err = m.AddWorkload(ctx, workload, nil)
    37  	assert.NoError(t, err)
    38  	// success updat
    39  	err = m.UpdateWorkload(ctx, workload)
    40  	assert.NoError(t, err)
    41  	// success updat
    42  	workload.Name = "test_app_2"
    43  	err = m.UpdateWorkload(ctx, workload)
    44  	assert.NoError(t, err)
    45  }
    46  
    47  func TestRemoveWorkload(t *testing.T) {
    48  	m := NewMercury(t)
    49  	ctx := context.Background()
    50  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
    51  	name := "test_app_1"
    52  	nodename := "n1"
    53  	podname := "test"
    54  	workload := &types.Workload{
    55  		ID:       ID,
    56  		Nodename: nodename,
    57  		Podname:  podname,
    58  		Name:     name,
    59  	}
    60  	// success create
    61  	err := m.AddWorkload(ctx, workload, nil)
    62  	assert.NoError(t, err)
    63  	// fail remove
    64  	workload.Name = "a"
    65  	err = m.RemoveWorkload(ctx, workload)
    66  	assert.Error(t, err)
    67  	workload.Name = name
    68  	// success remove
    69  	err = m.RemoveWorkload(ctx, workload)
    70  	assert.NoError(t, err)
    71  }
    72  
    73  func TestGetWorkload(t *testing.T) {
    74  	m := NewMercury(t)
    75  	ctx := context.Background()
    76  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
    77  	name := "test_app_1"
    78  	nodename := "n1"
    79  	podname := "test"
    80  	workload := &types.Workload{
    81  		ID:       ID,
    82  		Nodename: nodename,
    83  		Podname:  podname,
    84  		Name:     name,
    85  	}
    86  	// success create
    87  	err := m.AddWorkload(ctx, workload, nil)
    88  	assert.NoError(t, err)
    89  	// failed by no workload
    90  	_, err = m.GetWorkloads(ctx, []string{ID, "xxx"})
    91  	assert.Error(t, err)
    92  	// failed by no pod nodes
    93  	_, err = m.GetWorkload(ctx, ID)
    94  	assert.Error(t, err)
    95  	// create pod node
    96  	_, err = m.AddPod(ctx, podname, "")
    97  	assert.NoError(t, err)
    98  	_, err = m.AddNode(ctx, &types.AddNodeOptions{Nodename: nodename, Endpoint: "mock://", Podname: podname})
    99  	assert.NoError(t, err)
   100  	// success
   101  	_, err = m.GetWorkload(ctx, ID)
   102  	assert.NoError(t, err)
   103  }
   104  
   105  func TestGetWorkloadStatus(t *testing.T) {
   106  	m := NewMercury(t)
   107  	ctx := context.Background()
   108  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
   109  	name := "test_app_1"
   110  	nodename := "n1"
   111  	podname := "test"
   112  	workload := &types.Workload{
   113  		ID:       ID,
   114  		Nodename: nodename,
   115  		Podname:  podname,
   116  		Name:     name,
   117  	}
   118  	// success create
   119  	err := m.AddWorkload(ctx, workload, nil)
   120  	assert.NoError(t, err)
   121  	// failed no pod no node
   122  	_, err = m.GetWorkloadStatus(ctx, ID)
   123  	assert.Error(t, err)
   124  	// add success
   125  	_, err = m.AddPod(ctx, podname, "")
   126  	assert.NoError(t, err)
   127  	_, err = m.AddNode(ctx, &types.AddNodeOptions{Nodename: nodename, Endpoint: "mock://", Podname: podname})
   128  	assert.NoError(t, err)
   129  	c, err := m.GetWorkloadStatus(ctx, ID)
   130  	assert.Nil(t, c)
   131  }
   132  
   133  func TestSetWorkloadStatus(t *testing.T) {
   134  	m := NewMercury(t)
   135  	ctx := context.Background()
   136  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
   137  	name := "test_app_1"
   138  	nodename := "n1"
   139  	podname := "test"
   140  	workload := &types.Workload{
   141  		ID:         ID,
   142  		Nodename:   nodename,
   143  		Podname:    podname,
   144  		StatusMeta: &types.StatusMeta{ID: ID},
   145  	}
   146  	// fail by no name
   147  	err := m.SetWorkloadStatus(ctx, workload.StatusMeta, 0)
   148  	assert.Error(t, err)
   149  
   150  	workload.Name = name
   151  	workload.StatusMeta.Appname = "test"
   152  	workload.StatusMeta.Entrypoint = "app"
   153  	workload.StatusMeta.Nodename = "n1"
   154  	// no workload, err nil
   155  	err = m.SetWorkloadStatus(ctx, workload.StatusMeta, 10)
   156  	assert.Equal(t, err, types.ErrInvaildCount)
   157  	assert.NoError(t, m.AddWorkload(ctx, workload, nil))
   158  	// no status key, put succ, err nil
   159  	err = m.SetWorkloadStatus(ctx, workload.StatusMeta, 10)
   160  	assert.NoError(t, err)
   161  	// status not changed, update old lease
   162  	err = m.SetWorkloadStatus(ctx, workload.StatusMeta, 10)
   163  	assert.NoError(t, err)
   164  	// status changed, revoke old lease
   165  	workload.StatusMeta.Running = true
   166  	err = m.SetWorkloadStatus(ctx, workload.StatusMeta, 10)
   167  	assert.NoError(t, err)
   168  	// status not changed, ttl = 0
   169  	err = m.SetWorkloadStatus(ctx, workload.StatusMeta, 0)
   170  	assert.NoError(t, err)
   171  }
   172  
   173  func TestListWorkloads(t *testing.T) {
   174  	m := NewMercury(t)
   175  	ctx := context.Background()
   176  	// no key
   177  	cs, err := m.ListWorkloads(ctx, "", "a", "b", 1, nil)
   178  	assert.NoError(t, err)
   179  	assert.Empty(t, cs)
   180  	// add workload
   181  	name := "test_app_1"
   182  	nodename := "n1"
   183  	podname := "test"
   184  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
   185  	workload := &types.Workload{
   186  		ID:       ID,
   187  		Nodename: nodename,
   188  		Podname:  podname,
   189  		Name:     name,
   190  		Labels:   map[string]string{"x": "y"},
   191  	}
   192  	// success create
   193  	err = m.AddWorkload(ctx, workload, nil)
   194  	assert.NoError(t, err)
   195  	_, err = m.AddPod(ctx, podname, "")
   196  	assert.NoError(t, err)
   197  	_, err = m.AddNode(ctx, &types.AddNodeOptions{Nodename: nodename, Endpoint: "mock://", Podname: podname})
   198  	assert.NoError(t, err)
   199  	// no labels
   200  	cs, err = m.ListWorkloads(ctx, "", "a", "b", 1, nil)
   201  	assert.NoError(t, err)
   202  	assert.NotEmpty(t, cs)
   203  	// labels
   204  	cs, err = m.ListWorkloads(ctx, "", "a", "b", 1, map[string]string{"x": "z"})
   205  	assert.NoError(t, err)
   206  	assert.Empty(t, cs)
   207  }
   208  
   209  func TestListNodeWorkloads(t *testing.T) {
   210  	m := NewMercury(t)
   211  	ctx := context.Background()
   212  	// no key
   213  	cs, err := m.ListNodeWorkloads(ctx, "", nil)
   214  	assert.NoError(t, err)
   215  	assert.Empty(t, cs)
   216  	// add workload
   217  	name := "test_app_1"
   218  	nodename := "n1"
   219  	podname := "test"
   220  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
   221  	workload := &types.Workload{
   222  		ID:       ID,
   223  		Nodename: nodename,
   224  		Podname:  podname,
   225  		Name:     name,
   226  		Labels:   map[string]string{"x": "y"},
   227  	}
   228  	// success create
   229  	err = m.AddWorkload(ctx, workload, nil)
   230  	assert.NoError(t, err)
   231  	_, err = m.AddPod(ctx, podname, "")
   232  	assert.NoError(t, err)
   233  	_, err = m.AddNode(ctx, &types.AddNodeOptions{Nodename: nodename, Endpoint: "mock://", Podname: podname})
   234  	assert.NoError(t, err)
   235  	// no labels
   236  	cs, err = m.ListNodeWorkloads(ctx, nodename, nil)
   237  	assert.NoError(t, err)
   238  	assert.NotEmpty(t, cs)
   239  	// labels
   240  	cs, err = m.ListNodeWorkloads(ctx, nodename, map[string]string{"x": "z"})
   241  	assert.NoError(t, err)
   242  	assert.Empty(t, cs)
   243  }
   244  
   245  func TestWorkloadStatusStream(t *testing.T) {
   246  	m := NewMercury(t)
   247  	ctx := context.Background()
   248  	ID := "1234567812345678123456781234567812345678123456781234567812345678"
   249  	name := "test_app_1"
   250  	appname := "test"
   251  	entrypoint := "app"
   252  	nodename := "n1"
   253  	podname := "test"
   254  	workload := &types.Workload{
   255  		ID:         ID,
   256  		Name:       name,
   257  		Nodename:   nodename,
   258  		Podname:    podname,
   259  		StatusMeta: &types.StatusMeta{ID: ID},
   260  	}
   261  	node := &types.Node{
   262  		NodeMeta: types.NodeMeta{
   263  			Name:     nodename,
   264  			Podname:  podname,
   265  			Endpoint: "tcp://127.0.0.1:2376",
   266  		},
   267  	}
   268  	_, err := json.Marshal(workload)
   269  	assert.NoError(t, err)
   270  	nodeBytes, err := json.Marshal(node)
   271  	assert.NoError(t, err)
   272  	_, err = m.AddPod(ctx, podname, "CPU")
   273  	assert.NoError(t, err)
   274  	_, err = m.Create(ctx, fmt.Sprintf(nodeInfoKey, nodename), string(nodeBytes))
   275  	assert.NoError(t, err)
   276  	_, err = m.Create(ctx, fmt.Sprintf(nodePodKey, podname, nodename), string(nodeBytes))
   277  	assert.NoError(t, err)
   278  	assert.NoError(t, m.AddWorkload(ctx, workload, nil))
   279  	// WorkloadStatusStream
   280  	workload.StatusMeta = &types.StatusMeta{
   281  		ID:         ID,
   282  		Running:    true,
   283  		Appname:    appname,
   284  		Nodename:   nodename,
   285  		Entrypoint: entrypoint,
   286  	}
   287  	cctx, cancel := context.WithCancel(ctx)
   288  	ch := m.WorkloadStatusStream(cctx, appname, entrypoint, "", nil)
   289  	assert.NoError(t, m.SetWorkloadStatus(ctx, workload.StatusMeta, 0))
   290  	go func() {
   291  		time.Sleep(1 * time.Second)
   292  		cancel()
   293  	}()
   294  	for s := range ch {
   295  		assert.False(t, s.Delete)
   296  		assert.NotNil(t, s.Workload)
   297  	}
   298  }