github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/chunks/test_utils.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 2016 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 chunks
    23  
    24  import (
    25  	"context"
    26  	"sync/atomic"
    27  
    28  	"github.com/stretchr/testify/assert"
    29  
    30  	"github.com/dolthub/dolt/go/store/d"
    31  	"github.com/dolthub/dolt/go/store/hash"
    32  )
    33  
    34  func assertInputInStore(input string, h hash.Hash, s ChunkStore, assert *assert.Assertions) {
    35  	chunk, err := s.Get(context.Background(), h)
    36  	assert.NoError(err)
    37  	assert.False(chunk.IsEmpty(), "Shouldn't get empty chunk for %s", h.String())
    38  	assert.Equal(input, string(chunk.Data()))
    39  }
    40  
    41  func assertInputNotInStore(input string, h hash.Hash, s ChunkStore, assert *assert.Assertions) {
    42  	chunk, err := s.Get(context.Background(), h)
    43  	assert.NoError(err)
    44  	assert.True(chunk.IsEmpty(), "Shouldn't get non-empty chunk for %s: %v", h.String(), chunk)
    45  }
    46  
    47  type TestStorage struct {
    48  	MemoryStorage
    49  }
    50  
    51  func (t *TestStorage) NewView() *TestStoreView {
    52  	return &TestStoreView{ChunkStore: t.MemoryStorage.NewView()}
    53  }
    54  
    55  type TestStoreView struct {
    56  	ChunkStore
    57  	reads  int32
    58  	hases  int32
    59  	writes int32
    60  }
    61  
    62  var _ ChunkStoreGarbageCollector = &TestStoreView{}
    63  
    64  func (s *TestStoreView) Get(ctx context.Context, h hash.Hash) (Chunk, error) {
    65  	atomic.AddInt32(&s.reads, 1)
    66  	return s.ChunkStore.Get(ctx, h)
    67  }
    68  
    69  func (s *TestStoreView) GetMany(ctx context.Context, hashes hash.HashSet, found func(*Chunk)) error {
    70  	atomic.AddInt32(&s.reads, int32(len(hashes)))
    71  	return s.ChunkStore.GetMany(ctx, hashes, found)
    72  }
    73  
    74  func (s *TestStoreView) Has(ctx context.Context, h hash.Hash) (bool, error) {
    75  	atomic.AddInt32(&s.hases, 1)
    76  	return s.ChunkStore.Has(ctx, h)
    77  }
    78  
    79  func (s *TestStoreView) HasMany(ctx context.Context, hashes hash.HashSet) (hash.HashSet, error) {
    80  	atomic.AddInt32(&s.hases, int32(len(hashes)))
    81  	return s.ChunkStore.HasMany(ctx, hashes)
    82  }
    83  
    84  func (s *TestStoreView) Put(ctx context.Context, c Chunk) error {
    85  	atomic.AddInt32(&s.writes, 1)
    86  	return s.ChunkStore.Put(ctx, c)
    87  }
    88  
    89  func (s *TestStoreView) MarkAndSweepChunks(ctx context.Context, last hash.Hash, keepChunks <-chan []hash.Hash) error {
    90  	collector, ok := s.ChunkStore.(ChunkStoreGarbageCollector)
    91  	if !ok {
    92  		return ErrUnsupportedOperation
    93  	}
    94  
    95  	return collector.MarkAndSweepChunks(ctx, last, keepChunks)
    96  }
    97  
    98  func (s *TestStoreView) Reads() int {
    99  	reads := atomic.LoadInt32(&s.reads)
   100  	return int(reads)
   101  }
   102  
   103  func (s *TestStoreView) Hases() int {
   104  	hases := atomic.LoadInt32(&s.hases)
   105  	return int(hases)
   106  }
   107  
   108  func (s *TestStoreView) Writes() int {
   109  	writes := atomic.LoadInt32(&s.writes)
   110  	return int(writes)
   111  }
   112  
   113  type TestStoreFactory struct {
   114  	stores map[string]*TestStorage
   115  }
   116  
   117  func NewTestStoreFactory() *TestStoreFactory {
   118  	return &TestStoreFactory{map[string]*TestStorage{}}
   119  }
   120  
   121  func (f *TestStoreFactory) CreateStore(ns string) ChunkStore {
   122  	if f.stores == nil {
   123  		d.Panic("Cannot use TestStoreFactory after Shutter().")
   124  	}
   125  	if ts, present := f.stores[ns]; present {
   126  		return ts.NewView()
   127  	}
   128  	f.stores[ns] = &TestStorage{}
   129  	return f.stores[ns].NewView()
   130  }
   131  
   132  func (f *TestStoreFactory) Shutter() {
   133  	f.stores = nil
   134  }