github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/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/onflow/flow-go/engine/execution"
    12  	"github.com/onflow/flow-go/engine/execution/ingestion/uploader/mock"
    13  	executionUnittest "github.com/onflow/flow-go/engine/execution/state/unittest"
    14  	"github.com/onflow/flow-go/model/flow"
    15  	"github.com/onflow/flow-go/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(
    31  		t,
    32  		flow.ZeroID,
    33  		[][]flow.Identifier{
    34  			{flow.ZeroID},
    35  			{flow.ZeroID},
    36  			{flow.ZeroID},
    37  		})
    38  
    39  	t.Run("uploads with no errors", func(t *testing.T) {
    40  		uploader1 := mock.NewUploader(t)
    41  		uploader1.On("Upload", result).Return(nil).Once()
    42  
    43  		uploader2 := mock.NewUploader(t)
    44  		uploader2.On("Upload", result).Return(nil).Once()
    45  
    46  		uploader3 := mock.NewUploader(t)
    47  		uploader3.On("Upload", result).Return(nil).Once()
    48  
    49  		uploadMgr := NewManager(trace.NewNoopTracer())
    50  		uploadMgr.AddUploader(uploader1)
    51  		uploadMgr.AddUploader(uploader2)
    52  		uploadMgr.AddUploader(uploader3)
    53  
    54  		err := uploadMgr.Upload(context.Background(), result)
    55  		assert.NoError(t, err)
    56  	})
    57  
    58  	t.Run("uploads and returns error", func(t *testing.T) {
    59  		uploader2Err := fmt.Errorf("uploader 2 error")
    60  
    61  		uploader1 := mock.NewUploader(t)
    62  		uploader1.On("Upload", result).Return(nil).Once()
    63  
    64  		uploader2 := mock.NewUploader(t)
    65  		uploader2.On("Upload", result).Return(uploader2Err).Once()
    66  
    67  		uploader3 := mock.NewUploader(t)
    68  		uploader3.On("Upload", result).Return(nil).Once()
    69  
    70  		uploadMgr := NewManager(trace.NewNoopTracer())
    71  		uploadMgr.AddUploader(uploader1)
    72  		uploadMgr.AddUploader(uploader2)
    73  		uploadMgr.AddUploader(uploader3)
    74  
    75  		err := uploadMgr.Upload(context.Background(), result)
    76  		assert.ErrorIs(t, err, uploader2Err)
    77  	})
    78  }
    79  
    80  func TestRetryableUploader(t *testing.T) {
    81  	testRetryableUploader := new(FakeRetryableUploader)
    82  
    83  	uploadMgr := NewManager(trace.NewNoopTracer())
    84  	uploadMgr.AddUploader(testRetryableUploader)
    85  
    86  	err := uploadMgr.RetryUploads()
    87  	assert.Nil(t, err)
    88  
    89  	require.True(t, testRetryableUploader.RetryUploadCalled())
    90  }
    91  
    92  // FakeRetryableUploader is one RetryableUploader for testing purposes.
    93  type FakeRetryableUploader struct {
    94  	RetryableUploaderWrapper
    95  	retryUploadCalled bool
    96  }
    97  
    98  func (f *FakeRetryableUploader) Upload(_ *execution.ComputationResult) error {
    99  	return nil
   100  }
   101  
   102  func (f *FakeRetryableUploader) RetryUpload() error {
   103  	f.retryUploadCalled = true
   104  	return nil
   105  }
   106  
   107  func (f *FakeRetryableUploader) RetryUploadCalled() bool {
   108  	return f.retryUploadCalled
   109  }