storage/mockstorage/mockstorage.go

package mockstorage

import (
	"errors"
	"sync"
	"time"
)

// Config defines the config for mock storage.
type Config struct {
	CustomFuncs *CustomFuncs
}

// Storage is the mock storage adapter.
type Storage struct {
	mu     sync.RWMutex
	data   map[string]Entry
	custom *CustomFuncs
}

// Entry struct to hold value and expiration time.
type Entry struct {
	Value []byte
	Exp   time.Time
}

// CustomFuncs allows injecting custom behaviors for testing.
type CustomFuncs struct {
	GetFunc    func(key string) ([]byte, error)
	SetFunc    func(key string, val []byte, exp time.Duration) error
	DeleteFunc func(key string) error
	ResetFunc  func() error
	CloseFunc  func() error
	ConnFunc   func() map[string]Entry
	KeysFunc   func() ([][]byte, error)
}

// New creates a new mock storage with optional configuration.
func New(config ...Config) *Storage {
	s := &Storage{
		data: make(map[string]Entry),
		custom: &CustomFuncs{
			GetFunc:    nil,
			SetFunc:    nil,
			DeleteFunc: nil,
			ResetFunc:  nil,
			CloseFunc:  nil,
			ConnFunc:   nil,
			KeysFunc:   nil,
		},
	}

	// If a config is provided and it has CustomFuncs, use them
	if len(config) > 0 && config[0].CustomFuncs != nil {
		s.custom = config[0].CustomFuncs
	}

	return s
}

// Get retrieves the value for a given key.
func (s *Storage) Get(key string) ([]byte, error) {
	if s.custom.GetFunc != nil {
		return s.custom.GetFunc(key)
	}

	s.mu.RLock()
	defer s.mu.RUnlock()

	e, ok := s.data[key]
	if !ok {
		return nil, errors.New("key not found")
	}
	if !e.Exp.IsZero() && time.Now().After(e.Exp) {
		delete(s.data, key)
		return nil, errors.New("key expired")
	}
	return e.Value, nil
}

// Set sets the value for a given key with an expiration time.
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
	if s.custom.SetFunc != nil {
		return s.custom.SetFunc(key, val, exp)
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	var expTime time.Time
	if exp > 0 {
		expTime = time.Now().Add(exp)
	}

	s.data[key] = Entry{Value: val, Exp: expTime}
	return nil
}

// Delete removes a key from the storage.
func (s *Storage) Delete(key string) error {
	if s.custom.DeleteFunc != nil {
		return s.custom.DeleteFunc(key)
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	delete(s.data, key)
	return nil
}

// Reset clears all keys from the storage.
func (s *Storage) Reset() error {
	if s.custom.ResetFunc != nil {
		return s.custom.ResetFunc()
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	s.data = make(map[string]Entry)
	return nil
}

// Close closes the storage (no-op for mock).
func (s *Storage) Close() error {
	if s.custom.CloseFunc != nil {
		return s.custom.CloseFunc()
	}

	// No resources to clean up in mock
	return nil
}

// Conn returns the internal data map (for testing purposes).
func (s *Storage) Conn() map[string]Entry {
	if s.custom.ConnFunc != nil {
		return s.custom.ConnFunc()
	}

	s.mu.RLock()
	defer s.mu.RUnlock()

	copyData := make(map[string]Entry)
	for k, v := range s.data {
		copyData[k] = Entry{Value: v.Value, Exp: v.Exp}
	}
	return copyData
}

// Keys returns all keys in the storage.
func (s *Storage) Keys() ([][]byte, error) {
	if s.custom.KeysFunc != nil {
		return s.custom.KeysFunc()
	}

	s.mu.RLock()
	defer s.mu.RUnlock()

	keys := make([][]byte, 0, len(s.data))
	for k := range s.data {
		keys = append(keys, []byte(k))
	}
	return keys, nil
}

// SetCustomFuncs allows setting custom function implementations.
func (s *Storage) SetCustomFuncs(custom *CustomFuncs) {
	s.custom = custom
}