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

     1  package calcium
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io"
     7  	"testing"
     8  
     9  	enginemocks "github.com/projecteru2/core/engine/mocks"
    10  	storemocks "github.com/projecteru2/core/store/mocks"
    11  	"github.com/projecteru2/core/types"
    12  
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/mock"
    15  )
    16  
    17  func TestLogStream(t *testing.T) {
    18  	c := NewTestCluster()
    19  	store := c.store.(*storemocks.Store)
    20  	engine := &enginemocks.API{}
    21  	ID := "test"
    22  	workload := &types.Workload{
    23  		ID:     ID,
    24  		Engine: engine,
    25  	}
    26  	ctx := context.Background()
    27  	opts := &types.LogStreamOptions{ID: ID}
    28  	// failed by GetWorkload
    29  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
    30  	ch, err := c.LogStream(ctx, opts)
    31  	assert.NoError(t, err)
    32  	for c := range ch {
    33  		assert.Equal(t, c.ID, ID)
    34  		assert.Empty(t, c.Data)
    35  	}
    36  	store.On("GetWorkload", mock.Anything, mock.Anything).Return(workload, nil)
    37  	// failed by VirtualizationLogs
    38  	engine.On("VirtualizationLogs", mock.Anything, mock.Anything).Return(nil, nil, types.ErrMockError).Once()
    39  	ch, err = c.LogStream(ctx, opts)
    40  	assert.NoError(t, err)
    41  	for c := range ch {
    42  		assert.Equal(t, c.ID, ID)
    43  		assert.Empty(t, c.Data)
    44  	}
    45  	reader := bytes.NewBufferString("aaaa\nbbbb\n")
    46  	engine.On("VirtualizationLogs", mock.Anything, mock.Anything).Return(io.NopCloser(reader), nil, nil)
    47  	// success
    48  	ch, err = c.LogStream(ctx, opts)
    49  	assert.NoError(t, err)
    50  	for c := range ch {
    51  		assert.Equal(t, c.ID, ID)
    52  		assert.NotEmpty(t, c.Data)
    53  	}
    54  }