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

     1  package calcium
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"os"
     7  	"testing"
     8  
     9  	enginemocks "github.com/projecteru2/core/engine/mocks"
    10  	lockmocks "github.com/projecteru2/core/lock/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 TestSend(t *testing.T) {
    19  	c := NewTestCluster()
    20  	ctx := context.Background()
    21  
    22  	// 这部分是在测试参数合法性
    23  	// failed by validating
    24  	_, err := c.Send(ctx, &types.SendOptions{IDs: []string{}, Files: []types.LinuxFile{{Content: []byte("xxx")}}})
    25  	assert.Error(t, err)
    26  	_, err = c.Send(ctx, &types.SendOptions{IDs: []string{"id"}})
    27  	assert.Error(t, err)
    28  
    29  	tmpfile, err := os.CreateTemp("", "example")
    30  	assert.NoError(t, err)
    31  	defer os.RemoveAll(tmpfile.Name())
    32  	defer tmpfile.Close()
    33  	opts := &types.SendOptions{
    34  		IDs: []string{"cid"},
    35  		Files: []types.LinuxFile{
    36  			{
    37  				Filename: "/tmp/1",
    38  				Content:  []byte{},
    39  			},
    40  		},
    41  	}
    42  	store := c.store.(*storemocks.Store)
    43  	lock := &lockmocks.DistributedLock{}
    44  	lock.On("Lock", mock.Anything).Return(ctx, nil)
    45  	lock.On("Unlock", mock.Anything).Return(nil)
    46  	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
    47  	// failed by GetWorkload
    48  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    49  	ch, err := c.Send(ctx, opts)
    50  	assert.NoError(t, err)
    51  	for r := range ch {
    52  		assert.Error(t, r.Error)
    53  	}
    54  	engine := &enginemocks.API{}
    55  	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(
    56  		[]*types.Workload{{ID: "cid", Engine: engine}}, nil,
    57  	)
    58  	// failed by engine
    59  	content, _ := io.ReadAll(tmpfile)
    60  	opts.Files[0].Content = content
    61  	engine.On("VirtualizationCopyChunkTo",
    62  		mock.Anything, mock.Anything, mock.Anything,
    63  		mock.Anything, mock.Anything, mock.Anything,
    64  		mock.Anything, mock.Anything,
    65  	).Return(types.ErrMockError).Once()
    66  	ch, err = c.Send(ctx, opts)
    67  	assert.NoError(t, err)
    68  	for r := range ch {
    69  		assert.Error(t, r.Error)
    70  	}
    71  	// success
    72  	engine.On("VirtualizationCopyChunkTo",
    73  		mock.Anything, mock.Anything, mock.Anything,
    74  		mock.Anything, mock.Anything, mock.Anything,
    75  		mock.Anything, mock.Anything,
    76  	).Return(nil)
    77  	ch, err = c.Send(ctx, opts)
    78  	assert.NoError(t, err)
    79  	for r := range ch {
    80  		assert.NoError(t, r.Error)
    81  		assert.Equal(t, r.ID, "cid")
    82  		assert.Equal(t, r.Path, "/tmp/1")
    83  	}
    84  }