github.com/koko1123/flow-go-1@v0.29.6/engine/execution/ingestion/uploader/manager_test.go (about)

     1  package uploader
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/koko1123/flow-go-1/engine/execution"
    12  	"github.com/koko1123/flow-go-1/engine/execution/ingestion/uploader/mock"
    13  	executionUnittest "github.com/koko1123/flow-go-1/engine/execution/state/unittest"
    14  	"github.com/koko1123/flow-go-1/model/flow"
    15  	"github.com/koko1123/flow-go-1/module/trace"
    16  )
    17  
    18  func TestManagerSetEnabled(t *testing.T) {
    19  	uploadMgr := NewManager(trace.NewNoopTracer())
    20  	assert.False(t, uploadMgr.Enabled())
    21  
    22  	uploadMgr.SetEnabled(true)
    23  	assert.True(t, uploadMgr.Enabled())
    24  
    25  	uploadMgr.SetEnabled(false)
    26  	assert.False(t, uploadMgr.Enabled())
    27  }
    28  
    29  func TestManagerUploadsWithAllUploaders(t *testing.T) {
    30  	result := executionUnittest.ComputationResultFixture([][]flow.Identifier{
    31  		{flow.ZeroID},
    32  		{flow.ZeroID},
    33  		{flow.ZeroID},
    34  	})
    35  
    36  	t.Run("uploads with no errors", func(t *testing.T) {
    37  		uploader1 := mock.NewUploader(t)
    38  		uploader1.On("Upload", result).Return(nil).Once()
    39  
    40  		uploader2 := mock.NewUploader(t)
    41  		uploader2.On("Upload", result).Return(nil).Once()
    42  
    43  		uploader3 := mock.NewUploader(t)
    44  		uploader3.On("Upload", result).Return(nil).Once()
    45  
    46  		uploadMgr := NewManager(trace.NewNoopTracer())
    47  		uploadMgr.AddUploader(uploader1)
    48  		uploadMgr.AddUploader(uploader2)
    49  		uploadMgr.AddUploader(uploader3)
    50  
    51  		err := uploadMgr.Upload(context.Background(), result)
    52  		assert.NoError(t, err)
    53  	})
    54  
    55  	t.Run("uploads and returns error", func(t *testing.T) {
    56  		uploader2Err := fmt.Errorf("uploader 2 error")
    57  
    58  		uploader1 := mock.NewUploader(t)
    59  		uploader1.On("Upload", result).Return(nil).Once()
    60  
    61  		uploader2 := mock.NewUploader(t)
    62  		uploader2.On("Upload", result).Return(uploader2Err).Once()
    63  
    64  		uploader3 := mock.NewUploader(t)
    65  		uploader3.On("Upload", result).Return(nil).Once()
    66  
    67  		uploadMgr := NewManager(trace.NewNoopTracer())
    68  		uploadMgr.AddUploader(uploader1)
    69  		uploadMgr.AddUploader(uploader2)
    70  		uploadMgr.AddUploader(uploader3)
    71  
    72  		err := uploadMgr.Upload(context.Background(), result)
    73  		assert.ErrorIs(t, err, uploader2Err)
    74  	})
    75  }
    76  
    77  func TestRetryableUploader(t *testing.T) {
    78  	testRetryableUploader := new(FakeRetryableUploader)
    79  
    80  	uploadMgr := NewManager(trace.NewNoopTracer())
    81  	uploadMgr.AddUploader(testRetryableUploader)
    82  
    83  	err := uploadMgr.RetryUploads()
    84  	assert.Nil(t, err)
    85  
    86  	require.True(t, testRetryableUploader.RetryUploadCalled())
    87  }
    88  
    89  // FakeRetryableUploader is one RetryableUploader for testing purposes.
    90  type FakeRetryableUploader struct {
    91  	RetryableUploaderWrapper
    92  	retryUploadCalled bool
    93  }
    94  
    95  func (f *FakeRetryableUploader) Upload(_ *execution.ComputationResult) error {
    96  	return nil
    97  }
    98  
    99  func (f *FakeRetryableUploader) RetryUpload() error {
   100  	f.retryUploadCalled = true
   101  	return nil
   102  }
   103  
   104  func (f *FakeRetryableUploader) RetryUploadCalled() bool {
   105  	return f.retryUploadCalled
   106  }