storj.io/uplink@v1.13.0/private/storage/streams/pieceupload/upload_test.go (about)

     1  // Copyright (C) 2023 Storj Labs, Inc.
     2  // See LICENSE for copying information.
     3  
     4  package pieceupload
     5  
     6  import (
     7  	"bytes"
     8  	"context"
     9  	"io"
    10  	"testing"
    11  
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  	"github.com/zeebo/errs"
    15  
    16  	"storj.io/common/pb"
    17  	"storj.io/common/storj"
    18  )
    19  
    20  var (
    21  	fakePrivateKey = mustNewPiecePrivateKey()
    22  )
    23  
    24  func TestUploadOne(t *testing.T) {
    25  	for _, tc := range []struct {
    26  		desc           string
    27  		failPuts       int
    28  		cancelLongTail bool
    29  		cancelUpload   bool
    30  		expectUploaded bool
    31  		expectNum      int
    32  		expectErr      string
    33  	}{
    34  		{
    35  			desc:           "first piece successful",
    36  			expectUploaded: true,
    37  			expectNum:      0,
    38  		},
    39  		{
    40  			desc:           "second piece successful",
    41  			failPuts:       1,
    42  			expectUploaded: true,
    43  			expectNum:      1,
    44  		},
    45  		{
    46  			desc:         "upload canceled",
    47  			cancelUpload: true,
    48  			expectErr:    "context canceled",
    49  		},
    50  		{
    51  			desc:           "long tail canceled",
    52  			cancelLongTail: true,
    53  			expectUploaded: false,
    54  		},
    55  		{
    56  			desc:      "manager fails to return next piece",
    57  			failPuts:  2,
    58  			expectErr: "piece limit exchange failed: oh no",
    59  		},
    60  	} {
    61  		t.Run(tc.desc, func(t *testing.T) {
    62  			uploadCtx, uploadCancel := context.WithCancel(context.Background())
    63  			t.Cleanup(uploadCancel)
    64  			longTailCtx, longTailCancel := context.WithCancel(context.Background())
    65  			t.Cleanup(longTailCancel)
    66  
    67  			if tc.cancelUpload {
    68  				uploadCancel()
    69  			}
    70  			if tc.cancelLongTail {
    71  				longTailCancel()
    72  			}
    73  
    74  			manager := newManagerWithExchanger(2, failExchange{})
    75  			putter := &fakePutter{t: t, failPuts: tc.failPuts}
    76  			uploaded, err := UploadOne(longTailCtx, uploadCtx, manager, putter, fakePrivateKey)
    77  			if tc.expectErr != "" {
    78  				require.EqualError(t, err, tc.expectErr)
    79  				return
    80  			}
    81  			require.NoError(t, err)
    82  			require.Equal(t, tc.expectUploaded, uploaded)
    83  			if tc.expectUploaded {
    84  				assertResults(t, manager, revision{0}, makeResult(piecenum{tc.expectNum}, revision{0}))
    85  			}
    86  		})
    87  	}
    88  }
    89  
    90  type fakePutter struct {
    91  	t        *testing.T
    92  	failPuts int
    93  }
    94  
    95  func (p *fakePutter) PutPiece(longTailCtx, uploadCtx context.Context, limit *pb.AddressedOrderLimit, privateKey storj.PiecePrivateKey, data io.ReadCloser) (*pb.PieceHash, *struct{}, error) {
    96  	assert.Equal(p.t, fakePrivateKey, privateKey, "private key was not passed correctly")
    97  
    98  	num := pieceReaderNum(data)
    99  	if p.failPuts > 0 {
   100  		p.failPuts--
   101  		return nil, nil, errs.New("put failed for piece: %d", num)
   102  	}
   103  
   104  	select {
   105  	case <-uploadCtx.Done():
   106  		return nil, nil, uploadCtx.Err()
   107  	case <-longTailCtx.Done():
   108  		return nil, nil, longTailCtx.Err()
   109  	default:
   110  		return hash(num), nil, nil
   111  	}
   112  }
   113  
   114  func mustNewPiecePrivateKey() storj.PiecePrivateKey {
   115  	pk, err := storj.PiecePrivateKeyFromBytes(bytes.Repeat([]byte{1}, 64))
   116  	if err != nil {
   117  		panic(err)
   118  	}
   119  	return pk
   120  }