github.com/ethersphere/bee/v2@v2.2.0/pkg/statestore/mock/store.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 mock
     6  
     7  import (
     8  	"encoding"
     9  	"encoding/json"
    10  	"strings"
    11  	"sync"
    12  
    13  	"github.com/ethersphere/bee/v2/pkg/storage"
    14  )
    15  
    16  var _ storage.StateStorer = (*store)(nil)
    17  
    18  type store struct {
    19  	store map[string][]byte
    20  	mtx   sync.RWMutex
    21  }
    22  
    23  func NewStateStore() storage.StateStorer {
    24  	s := &store{
    25  		store: make(map[string][]byte),
    26  	}
    27  
    28  	return s
    29  }
    30  
    31  func (s *store) Get(key string, i interface{}) (err error) {
    32  	s.mtx.RLock()
    33  	defer s.mtx.RUnlock()
    34  
    35  	data, ok := s.store[key]
    36  	if !ok {
    37  		return storage.ErrNotFound
    38  	}
    39  
    40  	if unmarshaler, ok := i.(encoding.BinaryUnmarshaler); ok {
    41  		return unmarshaler.UnmarshalBinary(data)
    42  	}
    43  
    44  	return json.Unmarshal(data, i)
    45  }
    46  
    47  func (s *store) Put(key string, i interface{}) (err error) {
    48  	s.mtx.Lock()
    49  	defer s.mtx.Unlock()
    50  
    51  	var bytes []byte
    52  	if marshaler, ok := i.(encoding.BinaryMarshaler); ok {
    53  		if bytes, err = marshaler.MarshalBinary(); err != nil {
    54  			return err
    55  		}
    56  	} else if bytes, err = json.Marshal(i); err != nil {
    57  		return err
    58  	}
    59  
    60  	s.store[key] = bytes
    61  	return nil
    62  }
    63  
    64  func (s *store) Delete(key string) (err error) {
    65  	s.mtx.Lock()
    66  	defer s.mtx.Unlock()
    67  
    68  	delete(s.store, key)
    69  	return nil
    70  }
    71  
    72  func (s *store) Iterate(prefix string, iterFunc storage.StateIterFunc) (err error) {
    73  	s.mtx.RLock()
    74  	defer s.mtx.RUnlock()
    75  
    76  	for k, v := range s.store {
    77  		if !strings.HasPrefix(k, prefix) {
    78  			continue
    79  		}
    80  
    81  		val := make([]byte, len(v))
    82  		copy(val, v)
    83  		stop, err := iterFunc([]byte(k), val)
    84  		if err != nil {
    85  			return err
    86  		}
    87  
    88  		if stop {
    89  			return nil
    90  		}
    91  	}
    92  	return nil
    93  }
    94  
    95  func (s *store) Close() (err error) {
    96  	return nil
    97  }