github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/core/mock/dbdriver_mock.go (about)

     1  // Package mock provides a variety of mock implementations used for testing.
     2  /*
     3   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package mock
     6  
     7  import (
     8  	"sort"
     9  	"strings"
    10  	"sync"
    11  
    12  	"github.com/NVIDIA/aistore/cmn/cos"
    13  	"github.com/NVIDIA/aistore/cmn/kvdb"
    14  	jsoniter "github.com/json-iterator/go"
    15  )
    16  
    17  type DBDriver struct {
    18  	mtx    sync.RWMutex
    19  	values map[string]string
    20  }
    21  
    22  // interface guard
    23  var _ kvdb.Driver = (*DBDriver)(nil)
    24  
    25  func NewDBDriver() kvdb.Driver { return &DBDriver{values: make(map[string]string)} }
    26  func (*DBDriver) Close() error { return nil }
    27  
    28  func (*DBDriver) makePath(collection, key string) string {
    29  	return collection + kvdb.CollectionSepa + key
    30  }
    31  
    32  func (bd *DBDriver) Set(collection, key string, object any) error {
    33  	b := cos.MustMarshal(object)
    34  	return bd.SetString(collection, key, string(b))
    35  }
    36  
    37  func (bd *DBDriver) Get(collection, key string, object any) error {
    38  	s, err := bd.GetString(collection, key)
    39  	if err != nil {
    40  		return err
    41  	}
    42  	return jsoniter.Unmarshal([]byte(s), object)
    43  }
    44  
    45  func (bd *DBDriver) SetString(collection, key, data string) error {
    46  	bd.mtx.Lock()
    47  	defer bd.mtx.Unlock()
    48  	name := bd.makePath(collection, key)
    49  	bd.values[name] = data
    50  	return nil
    51  }
    52  
    53  func (bd *DBDriver) GetString(collection, key string) (string, error) {
    54  	bd.mtx.RLock()
    55  	defer bd.mtx.RUnlock()
    56  	name := bd.makePath(collection, key)
    57  	value, ok := bd.values[name]
    58  	if !ok {
    59  		return "", cos.NewErrNotFound(nil, collection+" \""+key+"\"")
    60  	}
    61  	return value, nil
    62  }
    63  
    64  func (bd *DBDriver) Delete(collection, key string) error {
    65  	bd.mtx.Lock()
    66  	defer bd.mtx.Unlock()
    67  	name := bd.makePath(collection, key)
    68  	_, ok := bd.values[name]
    69  	if !ok {
    70  		return cos.NewErrNotFound(nil, collection+" \""+key+"\"")
    71  	}
    72  	delete(bd.values, name)
    73  	return nil
    74  }
    75  
    76  func (bd *DBDriver) List(collection, pattern string) ([]string, error) {
    77  	var (
    78  		keys   = make([]string, 0)
    79  		filter string
    80  	)
    81  	bd.mtx.RLock()
    82  	defer bd.mtx.RUnlock()
    83  	filter = bd.makePath(collection, pattern)
    84  	for k := range bd.values {
    85  		if strings.HasPrefix(k, filter) {
    86  			_, key := kvdb.ParsePath(k)
    87  			if key != "" {
    88  				keys = append(keys, k)
    89  			}
    90  		}
    91  	}
    92  	sort.Strings(keys)
    93  	return keys, nil
    94  }
    95  
    96  func (bd *DBDriver) DeleteCollection(collection string) error {
    97  	keys, err := bd.List(collection, "")
    98  	bd.mtx.Lock()
    99  	defer bd.mtx.Unlock()
   100  	if err != nil || len(keys) == 0 {
   101  		return err
   102  	}
   103  	for _, k := range keys {
   104  		delete(bd.values, k)
   105  	}
   106  	return nil
   107  }
   108  
   109  func (bd *DBDriver) GetAll(collection, pattern string) (map[string]string, error) {
   110  	var (
   111  		values = make(map[string]string)
   112  		filter string
   113  	)
   114  	bd.mtx.RLock()
   115  	defer bd.mtx.RUnlock()
   116  	filter = bd.makePath(collection, pattern)
   117  	for k, v := range bd.values {
   118  		if strings.HasPrefix(k, filter) {
   119  			_, key := kvdb.ParsePath(k)
   120  			if key != "" {
   121  				values[key] = v
   122  			}
   123  		}
   124  	}
   125  	return values, nil
   126  }