add tests for lambda

Author: tonicmuroq

Makefile

@@ -31,6 +31,7 @@ mock: deps
 	mockery --dir store --output store/mocks --name Store
 	mockery --dir cluster --output cluster/mocks --name Cluster
 	mockery --dir lock --output lock/mocks --name DistributedLock
+	mockery --dir rpc/gen/ --output rpc/mocks --name CoreRPC_RunAndWaitServer
 
 .ONESHELL:
 

rpc/mocks/CoreRPC_RunAndWaitServer.go

@@ -870,7 +870,9 @@ func (v *Vibranium) RunAndWait(stream pb.CoreRPC_RunAndWaitServer) error {
 				}
 			}
 		})
+		return nil
 	}
+
 	go func() {
 		runAndWait(func(ch <-chan *types.AttachWorkloadMessage) {
 			r, w := io.Pipe()

rpc/rpc.go

@@ -1,14 +1,14 @@
 package rpc
 
 import (
-	"testing"
-
 	"context"
+	"testing"
 
 	clustermock "github.com/projecteru2/core/cluster/mocks"
 	enginemock "github.com/projecteru2/core/engine/mocks"
 	enginetypes "github.com/projecteru2/core/engine/types"
 	pb "github.com/projecteru2/core/rpc/gen"
+	rpcmocks "github.com/projecteru2/core/rpc/mocks"
 	"github.com/projecteru2/core/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -70,3 +70,125 @@ func TestSetNodeTranform(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, 2, len(o.DeltaCPU))
 }
+
+func TestRunAndWaitSync(t *testing.T) {
+	v := newVibranium()
+
+	stream := &rpcmocks.CoreRPC_RunAndWaitServer{}
+	stream.On("Context").Return(context.Background())
+	stream.On("Recv").Return(&pb.RunAndWaitOptions{
+		DeployOptions: &pb.DeployOptions{
+			Name: "deploy",
+			Entrypoint: &pb.EntrypointOptions{
+				Name:    "entry",
+				Command: "ping",
+			},
+			Podname:      "pod",
+			Image:        "image",
+			OpenStdin:    false,
+			ResourceOpts: &pb.ResourceOptions{},
+		},
+		Cmd:   []byte("ping"),
+		Async: false,
+	}, nil)
+
+	rc := []*pb.AttachWorkloadMessage{}
+	streamSendMock := func(m *pb.AttachWorkloadMessage) error {
+		rc = append(rc, m)
+		return nil
+	}
+	stream.On("Send", mock.Anything).Return(streamSendMock)
+
+	runAndWait := func(_ context.Context, _ *types.DeployOptions, _ <-chan []byte) <-chan *types.AttachWorkloadMessage {
+		ch := make(chan *types.AttachWorkloadMessage)
+		go func() {
+			// first message to report workload id
+			ch <- &types.AttachWorkloadMessage{
+				WorkloadID:    "workloadidfortonic",
+				Data:          []byte(""),
+				StdStreamType: types.Stdout,
+			}
+			// second message to report output of process
+			ch <- &types.AttachWorkloadMessage{
+				WorkloadID:    "workloadidfortonic",
+				Data:          []byte("network not reachable"),
+				StdStreamType: types.Stdout,
+			}
+			close(ch)
+		}()
+		return ch
+	}
+	cluster := v.cluster.(*clustermock.Cluster)
+	cluster.On("RunAndWait", mock.Anything, mock.Anything, mock.Anything).Return(runAndWait, nil)
+
+	err := v.RunAndWait(stream)
+	assert.NoError(t, err)
+	assert.Equal(t, len(rc), 2)
+
+	m1 := rc[0]
+	assert.Equal(t, m1.WorkloadId, "workloadidfortonic")
+	assert.Equal(t, m1.Data, []byte(""))
+
+	m2 := rc[1]
+	assert.Equal(t, m2.WorkloadId, "workloadidfortonic")
+	assert.Equal(t, m2.Data, []byte("network not reachable"))
+}
+
+func TestRunAndWaitAsync(t *testing.T) {
+	v := newVibranium()
+
+	stream := &rpcmocks.CoreRPC_RunAndWaitServer{}
+	stream.On("Context").Return(context.Background())
+	stream.On("Recv").Return(&pb.RunAndWaitOptions{
+		DeployOptions: &pb.DeployOptions{
+			Name: "deploy",
+			Entrypoint: &pb.EntrypointOptions{
+				Name:    "entry",
+				Command: "ping",
+			},
+			Podname:      "pod",
+			Image:        "image",
+			OpenStdin:    false,
+			ResourceOpts: &pb.ResourceOptions{},
+		},
+		Cmd:   []byte("ping"),
+		Async: true,
+	}, nil)
+
+	rc := []*pb.AttachWorkloadMessage{}
+	streamSendMock := func(m *pb.AttachWorkloadMessage) error {
+		rc = append(rc, m)
+		return nil
+	}
+	stream.On("Send", mock.Anything).Return(streamSendMock)
+
+	runAndWait := func(_ context.Context, _ *types.DeployOptions, _ <-chan []byte) <-chan *types.AttachWorkloadMessage {
+		ch := make(chan *types.AttachWorkloadMessage)
+		go func() {
+			// first message to report workload id
+			ch <- &types.AttachWorkloadMessage{
+				WorkloadID:    "workloadidfortonic",
+				Data:          []byte(""),
+				StdStreamType: types.Stdout,
+			}
+			// second message to report output of process
+			ch <- &types.AttachWorkloadMessage{
+				WorkloadID:    "workloadidfortonic",
+				Data:          []byte("network not reachable"),
+				StdStreamType: types.Stdout,
+			}
+			close(ch)
+		}()
+		return ch
+	}
+	cluster := v.cluster.(*clustermock.Cluster)
+	cluster.On("RunAndWait", mock.Anything, mock.Anything, mock.Anything).Return(runAndWait, nil)
+
+	err := v.RunAndWait(stream)
+	assert.NoError(t, err)
+	assert.Equal(t, len(rc), 1)
+
+	m1 := rc[0]
+	assert.Equal(t, m1.WorkloadId, "workloadidfortonic")
+	assert.Equal(t, m1.Data, []byte(""))
+}