github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/cmp_chunk_table_writer_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 package nbs 16 17 import ( 18 "bytes" 19 "context" 20 "testing" 21 22 "github.com/stretchr/testify/assert" 23 "github.com/stretchr/testify/require" 24 "golang.org/x/sync/errgroup" 25 26 "github.com/dolthub/dolt/go/store/chunks" 27 "github.com/dolthub/dolt/go/store/hash" 28 ) 29 30 func TestCmpChunkTableWriter(t *testing.T) { 31 // Put some chunks in a table file and get the buffer back which contains the table file data 32 ctx := context.Background() 33 34 expectedId, buff, err := WriteChunks(testMDChunks) 35 require.NoError(t, err) 36 37 // Setup a TableReader to read compressed chunks out of 38 ti, err := parseTableIndex(buff) 39 require.NoError(t, err) 40 tr := newTableReader(ti, tableReaderAtFromBytes(buff), fileBlockSize) 41 42 hashes := make(hash.HashSet) 43 for _, chnk := range testMDChunks { 44 hashes.Insert(chnk.Hash()) 45 } 46 47 reqs := toGetRecords(hashes) 48 found := make([]CompressedChunk, 0) 49 50 eg, egCtx := errgroup.WithContext(ctx) 51 _, err = tr.getManyCompressed(egCtx, eg, reqs, func(c CompressedChunk) { found = append(found, c) }, &Stats{}) 52 require.NoError(t, err) 53 require.NoError(t, eg.Wait()) 54 55 // for all the chunks we find, write them using the compressed writer 56 tw, err := NewCmpChunkTableWriter("") 57 require.NoError(t, err) 58 for _, cmpChnk := range found { 59 err = tw.AddCmpChunk(cmpChnk) 60 require.NoError(t, err) 61 err = tw.AddCmpChunk(cmpChnk) 62 assert.Equal(t, err, ErrChunkAlreadyWritten) 63 } 64 65 id, err := tw.Finish() 66 require.NoError(t, err) 67 68 assert.Equal(t, expectedId, id) 69 70 output := bytes.NewBuffer(nil) 71 err = tw.Flush(output) 72 require.NoError(t, err) 73 74 outputBuff := output.Bytes() 75 outputTI, err := parseTableIndex(outputBuff) 76 require.NoError(t, err) 77 outputTR := newTableReader(outputTI, tableReaderAtFromBytes(buff), fileBlockSize) 78 79 compareContentsOfTables(t, ctx, hashes, tr, outputTR) 80 } 81 82 func compareContentsOfTables(t *testing.T, ctx context.Context, hashes hash.HashSet, expectedRd, actualRd tableReader) { 83 expected, err := readAllChunks(ctx, hashes, expectedRd) 84 require.NoError(t, err) 85 actual, err := readAllChunks(ctx, hashes, actualRd) 86 require.NoError(t, err) 87 88 assert.Equal(t, len(expected), len(actual)) 89 assert.Equal(t, expected, actual) 90 } 91 92 func readAllChunks(ctx context.Context, hashes hash.HashSet, reader tableReader) (map[hash.Hash][]byte, error) { 93 reqs := toGetRecords(hashes) 94 found := make([]*chunks.Chunk, 0) 95 eg, ctx := errgroup.WithContext(ctx) 96 _, err := reader.getMany(ctx, eg, reqs, func(c *chunks.Chunk) { found = append(found, c) }, &Stats{}) 97 if err != nil { 98 return nil, err 99 } 100 err = eg.Wait() 101 if err != nil { 102 return nil, err 103 } 104 105 hashToData := make(map[hash.Hash][]byte) 106 for _, c := range found { 107 hashToData[c.Hash()] = c.Data() 108 } 109 110 return hashToData, nil 111 }