storj.io/uplink@v1.13.0/private/storage/streams/pieceupload/common_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  	"encoding/binary"
    10  	"errors"
    11  	"io"
    12  
    13  	"github.com/zeebo/errs"
    14  
    15  	"storj.io/common/pb"
    16  	"storj.io/common/storj"
    17  )
    18  
    19  // These are to improve the readability of the test cases themselves. Integer
    20  // constants are hard to follow.
    21  type revision struct{ value byte }
    22  type piecenum struct{ value int }
    23  
    24  func newManager(pieceCount int) *Manager {
    25  	limits := makeLimits(pieceCount)
    26  	exchanger := &fakeExchanger{
    27  		limits: map[byte][]*pb.AddressedOrderLimit{0: limits},
    28  	}
    29  	return NewManager(exchanger, new(fakePieceReader), makeSegmentID(revision{0}), limits)
    30  }
    31  
    32  func newManagerWithExchanger(pieceCount int, exchanger LimitsExchanger) *Manager {
    33  	limits := makeLimits(pieceCount)
    34  	return NewManager(exchanger, new(fakePieceReader), makeSegmentID(revision{0}), limits)
    35  }
    36  
    37  type fakePieceReader struct{}
    38  
    39  func (f fakePieceReader) PieceReader(num int) io.Reader {
    40  	var data [8]byte
    41  	binary.BigEndian.PutUint64(data[:], uint64(num))
    42  	return bytes.NewReader(data[:])
    43  }
    44  
    45  func pieceReaderNum(r io.Reader) piecenum {
    46  	var data [8]byte
    47  	_, _ = io.ReadFull(r, data[:])
    48  	return piecenum{int(binary.BigEndian.Uint64(data[:]))}
    49  }
    50  
    51  type failExchange struct{}
    52  
    53  func (failExchange) ExchangeLimits(ctx context.Context, segmentID storj.SegmentID, pieceNumbers []int) (storj.SegmentID, []*pb.AddressedOrderLimit, error) {
    54  	return nil, nil, errors.New("oh no")
    55  }
    56  
    57  type fakeExchanger struct {
    58  	limits map[byte][]*pb.AddressedOrderLimit
    59  }
    60  
    61  func (f fakeExchanger) ExchangeLimits(ctx context.Context, segmentID storj.SegmentID, pieceNumbers []int) (storj.SegmentID, []*pb.AddressedOrderLimit, error) {
    62  	if len(segmentID) != 1 {
    63  		return nil, nil, errs.New("programmer error: test segment IDs should be a single byte")
    64  	}
    65  
    66  	rev := segmentID[0]
    67  
    68  	limits := append([]*pb.AddressedOrderLimit(nil), f.limits[rev]...)
    69  	if len(limits) == 0 {
    70  		return nil, nil, errs.New("programmer error: no limits for segment ID revision %d", rev)
    71  	}
    72  
    73  	rev++
    74  
    75  	for _, num := range pieceNumbers {
    76  		if num < 0 || num >= len(limits) {
    77  			return nil, nil, errs.New("programmer error: piece number %d outside the range [0-%d)", num, len(limits))
    78  		}
    79  		limits[num] = makeLimit(piecenum{num}, revision{rev})
    80  	}
    81  
    82  	f.limits[rev] = limits
    83  
    84  	return makeSegmentID(revision{rev}), limits, nil
    85  }
    86  
    87  func makeLimits(n int) []*pb.AddressedOrderLimit {
    88  	var limits []*pb.AddressedOrderLimit
    89  	for num := 0; num < n; num++ {
    90  		limits = append(limits, makeLimit(piecenum{num}, revision{0}))
    91  	}
    92  	return limits
    93  }
    94  
    95  func makeLimit(num piecenum, rev revision) *pb.AddressedOrderLimit {
    96  	return &pb.AddressedOrderLimit{
    97  		Limit: &pb.OrderLimit{
    98  			PieceId:       pieceID(num),
    99  			StorageNodeId: nodeID(num, rev),
   100  		},
   101  	}
   102  }
   103  
   104  func makeSegmentID(rev revision) storj.SegmentID {
   105  	return storj.SegmentID{rev.value}
   106  }
   107  
   108  func hash(num piecenum) *pb.PieceHash {
   109  	return &pb.PieceHash{PieceId: pieceID(num)}
   110  }
   111  
   112  func pieceID(num piecenum) pb.PieceID {
   113  	var id pb.PieceID
   114  	binary.LittleEndian.PutUint64(id[:], uint64(num.value))
   115  	return id
   116  }
   117  
   118  func nodeID(num piecenum, rev revision) pb.NodeID {
   119  	var id pb.NodeID
   120  	binary.LittleEndian.PutUint64(id[:8], uint64(num.value))
   121  	binary.LittleEndian.PutUint64(id[8:], uint64(rev.value))
   122  	return id
   123  }