github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/aws_chunk_source_test.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2017 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package nbs
    23  
    24  import (
    25  	"context"
    26  	"testing"
    27  
    28  	"github.com/stretchr/testify/assert"
    29  	"github.com/stretchr/testify/require"
    30  )
    31  
    32  func TestAWSChunkSource(t *testing.T) {
    33  	chunks := [][]byte{
    34  		[]byte("hello2"),
    35  		[]byte("goodbye2"),
    36  		[]byte("badbye2"),
    37  	}
    38  	tableData, h, err := buildTable(chunks)
    39  	require.NoError(t, err)
    40  
    41  	s3 := makeFakeS3(t)
    42  	ddb := makeFakeDDB(t)
    43  
    44  	s3or := &s3ObjectReader{s3, "bucket", nil, nil, ""}
    45  	dts := &ddbTableStore{ddb, "table", nil, nil}
    46  
    47  	makeSrc := func(chunkMax int, ic *indexCache) chunkSource {
    48  		cs, err := newAWSChunkSource(
    49  			context.Background(),
    50  			dts,
    51  			s3or,
    52  			awsLimits{itemMax: maxDynamoItemSize, chunkMax: uint32(chunkMax)},
    53  			h,
    54  			uint32(len(chunks)),
    55  			ic,
    56  			&Stats{},
    57  			func(bs []byte) (tableIndex, error) {
    58  				return parseTableIndex(bs)
    59  			},
    60  		)
    61  
    62  		require.NoError(t, err)
    63  
    64  		return cs
    65  	}
    66  
    67  	t.Run("Dynamo", func(t *testing.T) {
    68  		ddb.putData(fmtTableName(h), tableData)
    69  
    70  		t.Run("NoIndexCache", func(t *testing.T) {
    71  			src := makeSrc(len(chunks)+1, nil)
    72  			assertChunksInReader(chunks, src, assert.New(t))
    73  		})
    74  
    75  		t.Run("WithIndexCache", func(t *testing.T) {
    76  			assert := assert.New(t)
    77  			index, err := parseTableIndex(tableData)
    78  			require.NoError(t, err)
    79  			cache := newIndexCache(1024)
    80  			cache.put(h, index)
    81  
    82  			baseline := ddb.NumGets()
    83  			src := makeSrc(len(chunks)+1, cache)
    84  
    85  			// constructing the table reader shouldn't have resulted in any reads
    86  			assert.Zero(ddb.NumGets() - baseline)
    87  			assertChunksInReader(chunks, src, assert)
    88  		})
    89  	})
    90  
    91  	t.Run("S3", func(t *testing.T) {
    92  		s3.data[h.String()] = tableData
    93  
    94  		t.Run("NoIndexCache", func(t *testing.T) {
    95  			src := makeSrc(len(chunks)-1, nil)
    96  			assertChunksInReader(chunks, src, assert.New(t))
    97  		})
    98  
    99  		t.Run("WithIndexCache", func(t *testing.T) {
   100  			assert := assert.New(t)
   101  			index, err := parseTableIndex(tableData)
   102  			require.NoError(t, err)
   103  			cache := newIndexCache(1024)
   104  			cache.put(h, index)
   105  
   106  			baseline := s3.getCount
   107  			src := makeSrc(len(chunks)-1, cache)
   108  
   109  			// constructing the table reader shouldn't have resulted in any reads
   110  			assert.Zero(s3.getCount - baseline)
   111  			assertChunksInReader(chunks, src, assert)
   112  		})
   113  	})
   114  }