github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/query/storage/mock/storage.go (about)

     1  // Copyright (c) 2018 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package mock
    22  
    23  import (
    24  	"context"
    25  	"errors"
    26  	"sync"
    27  	"time"
    28  
    29  	"github.com/m3db/m3/src/query/block"
    30  	"github.com/m3db/m3/src/query/storage"
    31  	"github.com/m3db/m3/src/query/storage/m3/consolidators"
    32  	"github.com/m3db/m3/src/query/storage/m3/storagemetadata"
    33  )
    34  
    35  // Storage implements storage.Storage and provides methods to help
    36  // read what was written and set what to retrieve.
    37  type Storage interface {
    38  	storage.Storage
    39  
    40  	SetTypeResult(storage.Type)
    41  	SetErrorBehavior(storage.ErrorBehavior)
    42  	LastFetchOptions() *storage.FetchOptions
    43  	SetFetchResult(*storage.FetchResult, error)
    44  	SetFetchResults(...*storage.FetchResult)
    45  	SetSearchSeriesResult(*storage.SearchResults, error)
    46  	SetCompleteTagsResult(*consolidators.CompleteTagsResult, error)
    47  	SetWriteResult(error)
    48  	SetFetchBlocksResult(block.Result, error)
    49  	SetQueryStorageMetadataAttributesResult([]storagemetadata.Attributes, error)
    50  	SetCloseResult(error)
    51  	Writes() []*storage.WriteQuery
    52  }
    53  
    54  type mockStorage struct {
    55  	sync.RWMutex
    56  	typeResult struct {
    57  		result storage.Type
    58  	}
    59  	errorBehavior    storage.ErrorBehavior
    60  	lastFetchOptions *storage.FetchOptions
    61  	fetchResult      struct {
    62  		results []*storage.FetchResult
    63  		idx     int
    64  		err     error
    65  	}
    66  	fetchTagsResult struct {
    67  		result *storage.SearchResults
    68  		err    error
    69  	}
    70  	writeResult struct {
    71  		err error
    72  	}
    73  	fetchBlocksResult struct {
    74  		result block.Result
    75  		err    error
    76  	}
    77  	completeTagsResult struct {
    78  		result *consolidators.CompleteTagsResult
    79  		err    error
    80  	}
    81  	closeResult struct {
    82  		err error
    83  	}
    84  	queryStorageMetadataAttributesResult struct {
    85  		attrs []storagemetadata.Attributes
    86  		err   error
    87  	}
    88  	writes []*storage.WriteQuery
    89  }
    90  
    91  // NewMockStorage creates a new mock Storage instance.
    92  func NewMockStorage() Storage {
    93  	return &mockStorage{}
    94  }
    95  
    96  func (s *mockStorage) FetchCompressed(ctx context.Context,
    97  	query *storage.FetchQuery, options *storage.FetchOptions) (consolidators.MultiFetchResult, error) {
    98  	panic("implement me")
    99  }
   100  
   101  func (s *mockStorage) SetTypeResult(result storage.Type) {
   102  	s.Lock()
   103  	defer s.Unlock()
   104  	s.typeResult.result = result
   105  }
   106  
   107  func (s *mockStorage) SetErrorBehavior(b storage.ErrorBehavior) {
   108  	s.Lock()
   109  	defer s.Unlock()
   110  	s.errorBehavior = b
   111  }
   112  
   113  func (s *mockStorage) SetFetchResult(result *storage.FetchResult, err error) {
   114  	s.Lock()
   115  	defer s.Unlock()
   116  	s.fetchResult.results = []*storage.FetchResult{result}
   117  	s.fetchResult.err = err
   118  }
   119  
   120  func (s *mockStorage) SetFetchResults(results ...*storage.FetchResult) {
   121  	s.Lock()
   122  	defer s.Unlock()
   123  	s.fetchResult.results = append(
   124  		make([]*storage.FetchResult, 0, len(results)),
   125  		results...,
   126  	)
   127  }
   128  
   129  func (s *mockStorage) SetSearchSeriesResult(result *storage.SearchResults, err error) {
   130  	s.Lock()
   131  	defer s.Unlock()
   132  	s.fetchTagsResult.result = result
   133  	s.fetchTagsResult.err = err
   134  }
   135  
   136  func (s *mockStorage) SetWriteResult(err error) {
   137  	s.Lock()
   138  	defer s.Unlock()
   139  	s.writeResult.err = err
   140  }
   141  
   142  func (s *mockStorage) SetFetchBlocksResult(result block.Result, err error) {
   143  	s.Lock()
   144  	defer s.Unlock()
   145  	s.fetchBlocksResult.result = result
   146  	s.fetchBlocksResult.err = err
   147  }
   148  
   149  func (s *mockStorage) SetCompleteTagsResult(result *consolidators.CompleteTagsResult, err error) {
   150  	s.Lock()
   151  	defer s.Unlock()
   152  	s.completeTagsResult.result = result
   153  	s.completeTagsResult.err = err
   154  }
   155  
   156  func (s *mockStorage) SetQueryStorageMetadataAttributesResult(attrs []storagemetadata.Attributes, err error) {
   157  	s.Lock()
   158  	defer s.Unlock()
   159  	s.queryStorageMetadataAttributesResult.attrs = attrs
   160  	s.queryStorageMetadataAttributesResult.err = err
   161  }
   162  
   163  func (s *mockStorage) SetCloseResult(err error) {
   164  	s.Lock()
   165  	defer s.Unlock()
   166  	s.closeResult.err = err
   167  }
   168  
   169  func (s *mockStorage) Writes() []*storage.WriteQuery {
   170  	s.RLock()
   171  	defer s.RUnlock()
   172  	return s.writes
   173  }
   174  
   175  func (s *mockStorage) LastFetchOptions() *storage.FetchOptions {
   176  	s.RLock()
   177  	defer s.RUnlock()
   178  	return s.lastFetchOptions
   179  }
   180  
   181  func (s *mockStorage) Fetch(
   182  	ctx context.Context,
   183  	query *storage.FetchQuery,
   184  	opts *storage.FetchOptions,
   185  ) (*storage.FetchResult, error) {
   186  	s.Lock()
   187  	defer s.Unlock()
   188  	s.lastFetchOptions = opts
   189  	idx := s.fetchResult.idx
   190  	if idx >= len(s.fetchResult.results) {
   191  		idx = 0
   192  	}
   193  
   194  	s.fetchResult.idx = s.fetchResult.idx + 1
   195  	return s.fetchResult.results[idx], s.fetchResult.err
   196  }
   197  
   198  func (s *mockStorage) FetchProm(
   199  	ctx context.Context,
   200  	query *storage.FetchQuery,
   201  	opts *storage.FetchOptions,
   202  ) (storage.PromResult, error) {
   203  	return storage.PromResult{}, errors.New("not implemented")
   204  }
   205  
   206  func (s *mockStorage) FetchBlocks(
   207  	ctx context.Context,
   208  	query *storage.FetchQuery,
   209  	opts *storage.FetchOptions,
   210  ) (block.Result, error) {
   211  	s.RLock()
   212  	defer s.RUnlock()
   213  	s.lastFetchOptions = opts
   214  	return s.fetchBlocksResult.result, s.fetchBlocksResult.err
   215  }
   216  
   217  func (s *mockStorage) SearchSeries(
   218  	ctx context.Context,
   219  	query *storage.FetchQuery,
   220  	opts *storage.FetchOptions,
   221  ) (*storage.SearchResults, error) {
   222  	s.RLock()
   223  	defer s.RUnlock()
   224  	s.lastFetchOptions = opts
   225  	return s.fetchTagsResult.result, s.fetchTagsResult.err
   226  }
   227  
   228  func (s *mockStorage) CompleteTags(
   229  	ctx context.Context,
   230  	query *storage.CompleteTagsQuery,
   231  	opts *storage.FetchOptions,
   232  ) (*consolidators.CompleteTagsResult, error) {
   233  	s.RLock()
   234  	defer s.RUnlock()
   235  	s.lastFetchOptions = opts
   236  	return s.completeTagsResult.result, s.completeTagsResult.err
   237  }
   238  
   239  func (s *mockStorage) QueryStorageMetadataAttributes(
   240  	_ context.Context,
   241  	_, _ time.Time,
   242  	opts *storage.FetchOptions,
   243  ) ([]storagemetadata.Attributes, error) {
   244  	s.RLock()
   245  	defer s.RUnlock()
   246  	s.lastFetchOptions = opts
   247  	return s.queryStorageMetadataAttributesResult.attrs, s.queryStorageMetadataAttributesResult.err
   248  }
   249  
   250  func (s *mockStorage) Write(
   251  	ctx context.Context,
   252  	query *storage.WriteQuery,
   253  ) error {
   254  	s.Lock()
   255  	defer s.Unlock()
   256  	s.writes = append(s.writes, query)
   257  	return s.writeResult.err
   258  }
   259  
   260  func (s *mockStorage) Type() storage.Type {
   261  	s.RLock()
   262  	defer s.RUnlock()
   263  	return s.typeResult.result
   264  }
   265  
   266  func (s *mockStorage) Name() string {
   267  	return "mock"
   268  }
   269  
   270  func (s *mockStorage) ErrorBehavior() storage.ErrorBehavior {
   271  	s.RLock()
   272  	defer s.RUnlock()
   273  	return s.errorBehavior
   274  }
   275  
   276  func (s *mockStorage) Close() error {
   277  	s.RLock()
   278  	defer s.RUnlock()
   279  	return s.closeResult.err
   280  }