storj.io/uplink@v1.13.0/private/storage/streams/batchaggregator/aggregator_test.go (about)

     1  // Copyright (C) 2023 Storj Labs, Inc.
     2  // See LICENSE for copying information.
     3  
     4  package batchaggregator
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  	"github.com/zeebo/errs"
    14  
    15  	"storj.io/common/pb"
    16  	"storj.io/uplink/private/metaclient"
    17  )
    18  
    19  func TestAggregator(t *testing.T) {
    20  	items := []metaclient.BatchItem{
    21  		&metaclient.BeginSegmentParams{StreamID: []byte("A")},
    22  		&metaclient.BeginSegmentParams{StreamID: []byte("B")},
    23  		&metaclient.BeginSegmentParams{StreamID: []byte("C")},
    24  		&metaclient.BeginSegmentParams{StreamID: []byte("D")},
    25  	}
    26  
    27  	responses := []*pb.BatchResponseItem{
    28  		{Response: &pb.BatchResponseItem_SegmentBegin{SegmentBegin: &pb.BeginSegmentResponse{SegmentId: []byte("1")}}},
    29  		{Response: &pb.BatchResponseItem_SegmentBegin{SegmentBegin: &pb.BeginSegmentResponse{SegmentId: []byte("2")}}},
    30  		{Response: &pb.BatchResponseItem_SegmentBegin{SegmentBegin: &pb.BeginSegmentResponse{SegmentId: []byte("3")}}},
    31  		{Response: &pb.BatchResponseItem_SegmentBegin{SegmentBegin: &pb.BeginSegmentResponse{SegmentId: []byte("4")}}},
    32  	}
    33  
    34  	t.Run("Schedule does not flush", func(t *testing.T) {
    35  		batcher := new(fakeBatcher)
    36  
    37  		aggregator := New(batcher)
    38  		aggregator.Schedule(items[0])
    39  		aggregator.Schedule(items[1])
    40  		aggregator.Schedule(items[2])
    41  
    42  		assert.Len(t, batcher.items, 0)
    43  	})
    44  
    45  	t.Run("ExecuteAndFlush flushes and returns last response with nothing else scheduled", func(t *testing.T) {
    46  		batcher := new(fakeBatcher)
    47  		batcher.responses = responses[:1]
    48  
    49  		aggregator := New(batcher)
    50  
    51  		resp, err := aggregator.ScheduleAndFlush(context.Background(), items[0])
    52  		require.NoError(t, err)
    53  		assert.Equal(t, items[:1], batcher.items)
    54  		assert.Equal(t, metaclient.MakeBatchResponse(items[0].BatchItem(), responses[0]), *resp)
    55  	})
    56  
    57  	t.Run("ExecuteAndFlush flushes and returns last response with other items scheduled", func(t *testing.T) {
    58  		batcher := new(fakeBatcher)
    59  		batcher.responses = responses[:4]
    60  
    61  		aggregator := New(batcher)
    62  		aggregator.Schedule(items[0])
    63  		aggregator.Schedule(items[1])
    64  		aggregator.Schedule(items[2])
    65  
    66  		resp, err := aggregator.ScheduleAndFlush(context.Background(), items[3])
    67  		require.NoError(t, err)
    68  		assert.Equal(t, items[:4], batcher.items)
    69  		assert.Equal(t, metaclient.MakeBatchResponse(items[3].BatchItem(), responses[3]), *resp)
    70  	})
    71  
    72  	t.Run("ExecuteAndFlush returns batch error", func(t *testing.T) {
    73  		batcher := new(fakeBatcher)
    74  		batcher.err = errors.New("oh no")
    75  
    76  		aggregator := New(batcher)
    77  
    78  		resp, err := aggregator.ScheduleAndFlush(context.Background(), items[0])
    79  		assert.EqualError(t, err, "oh no")
    80  		assert.Nil(t, resp)
    81  	})
    82  
    83  	t.Run("ExecuteAndFlush fails if batch response is empty", func(t *testing.T) {
    84  		batcher := new(fakeBatcher)
    85  		aggregator := New(batcher)
    86  
    87  		resp, err := aggregator.ScheduleAndFlush(context.Background(), items[0])
    88  		assert.EqualError(t, err, "missing batch responses")
    89  		assert.Nil(t, resp)
    90  	})
    91  
    92  	t.Run("Flush does not issue batch with nothing scheduled", func(t *testing.T) {
    93  		batcher := new(fakeBatcher)
    94  		aggregator := New(batcher)
    95  
    96  		require.NoError(t, aggregator.Flush(context.Background()))
    97  		assert.Empty(t, batcher.items)
    98  	})
    99  
   100  	t.Run("Flush flushes", func(t *testing.T) {
   101  		batcher := new(fakeBatcher)
   102  		batcher.responses = responses[3:]
   103  
   104  		aggregator := New(batcher)
   105  		aggregator.Schedule(items[0])
   106  		aggregator.Schedule(items[1])
   107  		aggregator.Schedule(items[2])
   108  
   109  		err := aggregator.Flush(context.Background())
   110  		require.NoError(t, err)
   111  		assert.Equal(t, items[:3], batcher.items)
   112  	})
   113  }
   114  
   115  type fakeBatcher struct {
   116  	items     []metaclient.BatchItem
   117  	responses []*pb.BatchResponseItem
   118  	err       error
   119  }
   120  
   121  func (mi *fakeBatcher) Batch(ctx context.Context, items ...metaclient.BatchItem) ([]metaclient.BatchResponse, error) {
   122  	if len(items) == 0 {
   123  		return nil, errs.New("test/programmer error: batch should never be issued with no items")
   124  	}
   125  	mi.items = items
   126  	if mi.err != nil {
   127  		return nil, mi.err
   128  	}
   129  	var responses []metaclient.BatchResponse
   130  	for i := 0; i < len(mi.responses); i++ {
   131  		item := items[i]
   132  		response := mi.responses[i]
   133  		responses = append(responses, metaclient.MakeBatchResponse(item.BatchItem(), response))
   134  	}
   135  	return responses, nil
   136  }