github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/rpc/rpc_test.go (about)

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	clustermock "github.com/projecteru2/core/cluster/mocks"
     8  	enginemock "github.com/projecteru2/core/engine/mocks"
     9  	enginetypes "github.com/projecteru2/core/engine/types"
    10  	pb "github.com/projecteru2/core/rpc/gen"
    11  	rpcmocks "github.com/projecteru2/core/rpc/mocks"
    12  	"github.com/projecteru2/core/types"
    13  
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/mock"
    16  )
    17  
    18  func newVibranium() *Vibranium {
    19  	v := &Vibranium{
    20  		cluster: &clustermock.Cluster{},
    21  	}
    22  	return v
    23  }
    24  
    25  func TestAddPod(t *testing.T) {
    26  	v := newVibranium()
    27  	opts := &pb.AddPodOptions{}
    28  	cluster := v.cluster.(*clustermock.Cluster)
    29  	cluster.On("AddPod", mock.Anything, mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    30  	_, err := v.AddPod(context.Background(), opts)
    31  	assert.Error(t, err)
    32  	cluster.On("AddPod", mock.Anything, mock.Anything, mock.Anything).Return(&types.Pod{Name: "test", Desc: "test"}, nil)
    33  	_, err = v.AddPod(context.Background(), opts)
    34  	assert.NoError(t, err)
    35  }
    36  
    37  func TestAddNode(t *testing.T) {
    38  	v := newVibranium()
    39  	opts := &pb.AddNodeOptions{}
    40  	cluster := v.cluster.(*clustermock.Cluster)
    41  	cluster.On("AddNode", mock.Anything, mock.Anything,
    42  		mock.Anything, mock.Anything, mock.Anything, mock.Anything,
    43  		mock.Anything, mock.Anything, mock.Anything, mock.Anything,
    44  		mock.Anything, mock.Anything, mock.Anything, mock.Anything,
    45  	).Return(nil, types.ErrMockError).Once()
    46  	_, err := v.AddNode(context.Background(), opts)
    47  	assert.Error(t, err)
    48  	engine := &enginemock.API{}
    49  	node := &types.Node{
    50  		NodeMeta: types.NodeMeta{
    51  			Name: "test",
    52  		},
    53  		Engine: engine,
    54  	}
    55  	engine.On("Info", mock.Anything).Return(&enginetypes.Info{}, nil)
    56  	cluster.On("AddNode", mock.Anything, mock.Anything,
    57  		mock.Anything, mock.Anything, mock.Anything, mock.Anything,
    58  		mock.Anything, mock.Anything, mock.Anything, mock.Anything,
    59  		mock.Anything, mock.Anything, mock.Anything, mock.Anything,
    60  	).Return(node, nil)
    61  	_, err = v.AddNode(context.Background(), opts)
    62  	assert.NoError(t, err)
    63  }
    64  
    65  func TestSetNodeTranform(t *testing.T) {
    66  	b := &pb.SetNodeOptions{
    67  		Nodename: "a",
    68  	}
    69  	_, err := toCoreSetNodeOptions(b)
    70  	assert.Nil(t, err)
    71  }
    72  
    73  func TestRunAndWaitSync(t *testing.T) {
    74  	v := newVibranium()
    75  
    76  	stream := &rpcmocks.CoreRPC_RunAndWaitServer{}
    77  	stream.On("Context").Return(context.Background())
    78  	stream.On("Recv").Return(&pb.RunAndWaitOptions{
    79  		DeployOptions: &pb.DeployOptions{
    80  			Name: "deploy",
    81  			Entrypoint: &pb.EntrypointOptions{
    82  				Name:    "entry",
    83  				Command: "ping",
    84  			},
    85  			Podname:   "pod",
    86  			Image:     "image",
    87  			OpenStdin: false,
    88  			Resources: map[string][]byte{},
    89  		},
    90  		Cmd:   []byte("ping"),
    91  		Async: false,
    92  	}, nil)
    93  
    94  	rc := []*pb.AttachWorkloadMessage{}
    95  	streamSendMock := func(m *pb.AttachWorkloadMessage) error {
    96  		rc = append(rc, m)
    97  		return nil
    98  	}
    99  	stream.On("Send", mock.Anything).Return(streamSendMock)
   100  
   101  	runAndWait := func(_ context.Context, _ *types.DeployOptions, _ <-chan []byte) <-chan *types.AttachWorkloadMessage {
   102  		ch := make(chan *types.AttachWorkloadMessage)
   103  		go func() {
   104  			// message to report output of process
   105  			ch <- &types.AttachWorkloadMessage{
   106  				WorkloadID:    "workloadidfortonic",
   107  				Data:          []byte("network not reachable"),
   108  				StdStreamType: types.Stderr,
   109  			}
   110  			close(ch)
   111  		}()
   112  		return ch
   113  	}
   114  	cluster := v.cluster.(*clustermock.Cluster)
   115  	cluster.On("RunAndWait", mock.Anything, mock.Anything, mock.Anything).Return([]string{"workloadidfortonic"}, runAndWait, nil)
   116  
   117  	err := v.RunAndWait(stream)
   118  	assert.NoError(t, err)
   119  	assert.Equal(t, len(rc), 2)
   120  
   121  	m1 := rc[0]
   122  	assert.Equal(t, m1.WorkloadId, "workloadidfortonic")
   123  	assert.Equal(t, m1.Data, []byte(""))
   124  	assert.Equal(t, m1.StdStreamType, pb.StdStreamType_TYPEWORKLOADID)
   125  
   126  	m2 := rc[1]
   127  	assert.Equal(t, m2.WorkloadId, "workloadidfortonic")
   128  	assert.Equal(t, m2.Data, []byte("network not reachable"))
   129  	assert.Equal(t, m2.StdStreamType, pb.StdStreamType_STDERR)
   130  }
   131  
   132  func TestRunAndWaitAsync(t *testing.T) {
   133  	v := newVibranium()
   134  
   135  	stream := &rpcmocks.CoreRPC_RunAndWaitServer{}
   136  	stream.On("Context").Return(context.Background())
   137  	stream.On("Recv").Return(&pb.RunAndWaitOptions{
   138  		DeployOptions: &pb.DeployOptions{
   139  			Name: "deploy",
   140  			Entrypoint: &pb.EntrypointOptions{
   141  				Name:    "entry",
   142  				Command: "ping",
   143  			},
   144  			Podname:   "pod",
   145  			Image:     "image",
   146  			OpenStdin: false,
   147  			Resources: map[string][]byte{},
   148  		},
   149  		Cmd:   []byte("ping"),
   150  		Async: true,
   151  	}, nil)
   152  
   153  	rc := []*pb.AttachWorkloadMessage{}
   154  	streamSendMock := func(m *pb.AttachWorkloadMessage) error {
   155  		rc = append(rc, m)
   156  		return nil
   157  	}
   158  	stream.On("Send", mock.Anything).Return(streamSendMock)
   159  
   160  	runAndWait := func(_ context.Context, _ *types.DeployOptions, _ <-chan []byte) <-chan *types.AttachWorkloadMessage {
   161  		ch := make(chan *types.AttachWorkloadMessage)
   162  		go func() {
   163  			// message to report output of process
   164  			ch <- &types.AttachWorkloadMessage{
   165  				WorkloadID:    "workloadidfortonic",
   166  				Data:          []byte("network not reachable"),
   167  				StdStreamType: types.Stderr,
   168  			}
   169  			close(ch)
   170  		}()
   171  		return ch
   172  	}
   173  	cluster := v.cluster.(*clustermock.Cluster)
   174  	cluster.On("RunAndWait", mock.Anything, mock.Anything, mock.Anything).Return([]string{"workloadidfortonic"}, runAndWait, nil)
   175  
   176  	err := v.RunAndWait(stream)
   177  	assert.NoError(t, err)
   178  	assert.Equal(t, len(rc), 1)
   179  
   180  	m1 := rc[0]
   181  	assert.Equal(t, m1.WorkloadId, "workloadidfortonic")
   182  	assert.Equal(t, m1.Data, []byte(""))
   183  	assert.Equal(t, m1.StdStreamType, pb.StdStreamType_TYPEWORKLOADID)
   184  }