github.com/ethersphere/bee/v2@v2.2.0/pkg/postage/testing/batch.go (about)

     1  // Copyright 2020 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package testing
     6  
     7  import (
     8  	"bytes"
     9  	crand "crypto/rand"
    10  	"io"
    11  	"math/big"
    12  	"math/rand"
    13  	"testing"
    14  
    15  	"github.com/ethersphere/bee/v2/pkg/postage"
    16  )
    17  
    18  const (
    19  	defaultBucketDepth = 12
    20  	defaultDepth       = 16
    21  )
    22  
    23  // BatchOption is an optional parameter for NewBatch
    24  type BatchOption func(c *postage.Batch)
    25  
    26  // MustNewID will generate a new random ID (32 byte slice). Panics on errors.
    27  func MustNewID() []byte {
    28  	id := make([]byte, 32)
    29  	_, err := io.ReadFull(crand.Reader, id)
    30  	if err != nil {
    31  		panic(err)
    32  	}
    33  	return id
    34  }
    35  
    36  // MustNewAddress will generate a new random address (20 byte slice). Panics on
    37  // errors.
    38  func MustNewAddress() []byte {
    39  	addr := make([]byte, 20)
    40  	_, err := io.ReadFull(crand.Reader, addr)
    41  	if err != nil {
    42  		panic(err)
    43  	}
    44  	return addr
    45  }
    46  
    47  // NewBigInt will generate a new random big int (uint64 base value).
    48  func NewBigInt() *big.Int {
    49  	return (new(big.Int)).SetUint64(rand.Uint64())
    50  }
    51  
    52  // MustNewBatch will create a new test batch. Fields that are not supplied will
    53  // be filled with random data. Panics on errors.
    54  func MustNewBatch(opts ...BatchOption) *postage.Batch {
    55  	b := &postage.Batch{
    56  		ID:          MustNewID(),
    57  		Value:       NewBigInt(),
    58  		Start:       rand.Uint64(),
    59  		BucketDepth: defaultBucketDepth,
    60  		Depth:       defaultDepth,
    61  		Immutable:   true,
    62  	}
    63  
    64  	for _, opt := range opts {
    65  		opt(b)
    66  	}
    67  
    68  	if b.Owner == nil {
    69  		b.Owner = MustNewAddress()
    70  	}
    71  
    72  	return b
    73  }
    74  
    75  // WithOwner will set the batch owner on a randomized batch.
    76  func WithOwner(owner []byte) BatchOption {
    77  	return func(b *postage.Batch) {
    78  		b.Owner = owner
    79  	}
    80  }
    81  
    82  // WithValue will set the batch Value to the given value.
    83  func WithValue(value int64) BatchOption {
    84  	return func(b *postage.Batch) {
    85  		b.Value = big.NewInt(value)
    86  	}
    87  }
    88  
    89  // WithDepth will set the batch Depth to the given depth.
    90  func WithDepth(depth uint8) BatchOption {
    91  	return func(b *postage.Batch) {
    92  		b.Depth = depth
    93  	}
    94  }
    95  
    96  // WithStart will set the batch Start to the given start.
    97  func WithStart(start uint64) BatchOption {
    98  	return func(b *postage.Batch) {
    99  		b.Start = start
   100  	}
   101  }
   102  
   103  // CompareBatches is a testing helper that compares two batches and fails the
   104  // test if all fields are not equal.
   105  // Fails on first different value and prints the comparison.
   106  func CompareBatches(t *testing.T, want, got *postage.Batch) {
   107  	t.Helper()
   108  
   109  	if !bytes.Equal(want.ID, got.ID) {
   110  		t.Fatalf("batch ID: want %v, got %v", want.ID, got.ID)
   111  	}
   112  	if want.Value.Cmp(got.Value) != 0 {
   113  		t.Fatalf("value: want %v, got %v", want.Value, got.Value)
   114  	}
   115  	if want.Start != got.Start {
   116  		t.Fatalf("start: want %v, got %b", want.Start, got.Start)
   117  	}
   118  	if !bytes.Equal(want.Owner, got.Owner) {
   119  		t.Fatalf("owner: want %v, got %v", want.Owner, got.Owner)
   120  	}
   121  	if want.Depth != got.Depth {
   122  		t.Fatalf("depth: want %v, got %v", want.Depth, got.Depth)
   123  	}
   124  }