navidrome/tests/mock_artwork_queue_repo.go

package tests

import (
	"cmp"
	"slices"
	"sync"
	"time"

	"github.com/navidrome/navidrome/model"
	"github.com/navidrome/navidrome/utils/slice"
)

type MockArtworkQueueRepo struct {
	model.ArtworkQueueRepository
	// mu guards Data: the worker drains this mock concurrently.
	mu   sync.Mutex
	Data map[string]model.ArtworkQueueItem // keyed by iaKey(kind, id, imageType)
	Err  error
	// ItemArtworkSource, when set, backs EnqueueStaleAbsent with real item_artwork state.
	ItemArtworkSource *MockArtworkRepo
	// ExistingIDs is keyed by item_kind; a nil per-kind map means PurgeDangling keeps that kind.
	ExistingIDs map[string]map[string]bool
}

func CreateMockArtworkQueueRepo() *MockArtworkQueueRepo {
	return &MockArtworkQueueRepo{Data: map[string]model.ArtworkQueueItem{}}
}

func (m *MockArtworkQueueRepo) Get(kind model.Kind, id, imageType string) (*model.ArtworkQueueItem, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return nil, m.Err
	}
	it, ok := m.Data[iaKey(kind.Prefix(), id, imageType)]
	if !ok {
		return nil, model.ErrNotFound
	}
	return &it, nil
}

func (m *MockArtworkQueueRepo) Enqueue(items ...model.ArtworkQueueItem) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return m.Err
	}
	m.enqueueLocked(items)
	return nil
}

func (m *MockArtworkQueueRepo) enqueueLocked(items []model.ArtworkQueueItem) {
	now := time.Now()
	for _, it := range items {
		if it.ImageType == "" {
			it.ImageType = model.ImageTypePrimary
		}
		k := iaKey(it.ItemKind, it.ItemID, it.ImageType)
		// Mirror the SQL: retry_at/enqueued_at are server-set, never taken from the caller.
		if prev, ok := m.Data[k]; ok {
			prev.Priority = max(prev.Priority, it.Priority)
			prev.RetryAt = now
			prev.Attempts = 0
			prev.EnqueuedAt = now
			m.Data[k] = prev
			continue
		}
		it.Attempts = 0
		it.RetryAt = now
		it.EnqueuedAt = now
		m.Data[k] = it
	}
}

// EnqueueIfMissing mirrors the SQL anti-join: skip anything that already has an item_artwork row.
func (m *MockArtworkQueueRepo) EnqueueIfMissing(items ...model.ArtworkQueueItem) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return m.Err
	}
	var fresh []model.ArtworkQueueItem
	for _, it := range items {
		k := iaKey(it.ItemKind, it.ItemID, cmp.Or(it.ImageType, model.ImageTypePrimary))
		if m.ItemArtworkSource != nil {
			if _, ok := m.ItemArtworkSource.ItemData[k]; ok {
				continue
			}
		}
		if _, ok := m.Data[k]; ok { // DO NOTHING
			continue
		}
		fresh = append(fresh, it)
	}
	m.enqueueLocked(fresh)
	return nil
}

func (m *MockArtworkQueueRepo) DequeueBatch(n int, kinds ...string) ([]model.ArtworkQueueItem, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return nil, m.Err
	}
	var res []model.ArtworkQueueItem
	now := time.Now()
	for _, it := range m.Data {
		if !it.RetryAt.After(now) && (len(kinds) == 0 || slices.Contains(kinds, it.ItemKind)) {
			res = append(res, it)
		}
	}
	slices.SortFunc(res, func(a, b model.ArtworkQueueItem) int {
		return cmp.Or(cmp.Compare(b.Priority, a.Priority), a.EnqueuedAt.Compare(b.EnqueuedAt))
	})
	if len(res) > n {
		res = res[:n]
	}
	return res, nil
}

func (m *MockArtworkQueueRepo) MarkFailedIfUnchanged(kind, id, imageType string, seenRetryAt, retryAt time.Time, trace string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return m.Err
	}
	k := iaKey(kind, id, imageType)
	if it, ok := m.Data[k]; ok && it.RetryAt.Equal(seenRetryAt) {
		it.Attempts++
		it.RetryAt = retryAt
		it.Trace = trace
		m.Data[k] = it
	}
	return nil
}

func (m *MockArtworkQueueRepo) DeleteIfUnchanged(kind, id, imageType string, retryAt time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return m.Err
	}
	k := iaKey(kind, id, imageType)
	if it, ok := m.Data[k]; ok && it.RetryAt.Equal(retryAt) {
		delete(m.Data, k)
	}
	return nil
}

func (m *MockArtworkQueueRepo) PurgeDangling() (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return 0, m.Err
	}
	var purged int64
	for k, it := range m.Data {
		existing := m.ExistingIDs[it.ItemKind]
		if existing == nil {
			continue
		}
		if !existing[it.ItemID] {
			delete(m.Data, k)
			purged++
		}
	}
	return purged, nil
}

// queueFilterMatches mirrors artworkQueueFilter, so the mock cannot let a preview and a delete disagree.
func queueFilterMatches(it model.ArtworkQueueItem, kinds []model.Kind, priorities []int) bool {
	prefixes := model.KindPrefixes(kinds)
	return (len(prefixes) == 0 || slices.Contains(prefixes, it.ItemKind)) &&
		(len(priorities) == 0 || slices.Contains(priorities, it.Priority))
}

func (m *MockArtworkQueueRepo) PurgeQueued(kinds []model.Kind, priorities []int) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return 0, m.Err
	}
	var purged int64
	for k, it := range m.Data {
		if !queueFilterMatches(it, kinds, priorities) {
			continue
		}
		delete(m.Data, k)
		purged++
	}
	return purged, nil
}

func (m *MockArtworkQueueRepo) Count() (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return 0, m.Err
	}
	return int64(len(m.Data)), nil
}

func (m *MockArtworkQueueRepo) CountQueued(kinds []model.Kind, priorities []int) ([]model.ArtworkQueueStat, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return nil, m.Err
	}
	var res []model.ArtworkQueueStat
	for _, it := range m.Data {
		if !queueFilterMatches(it, kinds, priorities) {
			continue
		}
		i := slices.IndexFunc(res, func(s model.ArtworkQueueStat) bool {
			return s.ItemKind == it.ItemKind && s.Priority == it.Priority
		})
		if i < 0 {
			res = append(res, model.ArtworkQueueStat{ItemKind: it.ItemKind, Priority: it.Priority, Count: 1})
			continue
		}
		res[i].Count++
	}
	slices.SortFunc(res, func(a, b model.ArtworkQueueStat) int {
		return cmp.Or(cmp.Compare(a.ItemKind, b.ItemKind), cmp.Compare(b.Priority, a.Priority))
	})
	return res, nil
}

// CountAbsent mirrors the SQL predicate: an absent state is one with no hash.
func (m *MockArtworkQueueRepo) CountAbsent(kind model.Kind, attemptedBefore time.Time) (model.ArtworkAbsentStat, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var res model.ArtworkAbsentStat
	if m.Err != nil || m.ItemArtworkSource == nil {
		return res, m.Err
	}
	for _, ia := range m.ItemArtworkSource.ItemData {
		if ia.ItemKind != kind.Prefix() || ia.Hash != "" {
			continue
		}
		res.Total++
		if ia.AttemptedAt.Before(attemptedBefore) {
			res.Stale++
		}
	}
	return res, nil
}

func (m *MockArtworkQueueRepo) EnqueuePreservingBackoff(items ...model.ArtworkQueueItem) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return m.Err
	}
	now := time.Now()
	for _, it := range items {
		if it.ImageType == "" {
			it.ImageType = model.ImageTypePrimary
		}
		k := iaKey(it.ItemKind, it.ItemID, it.ImageType)
		// Preserve an existing row's retry_at: a bump raises priority without resetting backoff.
		if prev, ok := m.Data[k]; ok {
			prev.Priority = max(prev.Priority, it.Priority)
			m.Data[k] = prev
			continue
		}
		it.Attempts = 0
		it.RetryAt = now
		it.EnqueuedAt = now
		m.Data[k] = it
	}
	return nil
}

func (m *MockArtworkQueueRepo) EnqueueStaleAbsent(kind model.Kind, attemptedBefore time.Time, limit int) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil || m.ItemArtworkSource == nil {
		return 0, m.Err
	}
	var stale []model.ItemArtwork
	for _, ia := range m.ItemArtworkSource.ItemData {
		if ia.ItemKind == kind.Prefix() && ia.Hash == "" && ia.AttemptedAt.Before(attemptedBefore) {
			stale = append(stale, ia)
		}
	}
	slices.SortFunc(stale, func(a, b model.ItemArtwork) int { return a.AttemptedAt.Compare(b.AttemptedAt) })
	// The limit caps the selection, like the SQL's LIMIT before ON CONFLICT: queued rows use up budget.
	stale = stale[:min(limit, len(stale))]
	now := time.Now()
	var inserted int64
	for _, ia := range stale {
		k := iaKey(ia.ItemKind, ia.ItemID, ia.ImageType)
		if _, ok := m.Data[k]; ok { // DO NOTHING: never touch existing queue rows
			continue
		}
		m.Data[k] = model.ArtworkQueueItem{
			ItemKind:   ia.ItemKind,
			ItemID:     ia.ItemID,
			ImageType:  ia.ImageType,
			Priority:   model.ArtworkPriorityRecheck,
			RetryAt:    now,
			EnqueuedAt: now,
		}
		inserted++
	}
	return inserted, nil
}

// matchingSource mirrors the SQL filter: no sources means every source, "" the absent state.
func (m *MockArtworkQueueRepo) matchingSource(kind model.Kind, sources []string) []model.ItemArtwork {
	if m.ItemArtworkSource == nil {
		return nil
	}
	var res []model.ItemArtwork
	for _, ia := range m.ItemArtworkSource.ItemData {
		if ia.ItemKind == kind.Prefix() && (len(sources) == 0 || slices.Contains(sources, ia.Source)) {
			res = append(res, ia)
		}
	}
	return res
}

func (m *MockArtworkQueueRepo) CountBySource(kind model.Kind, sources []string) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return 0, m.Err
	}
	return int64(len(m.matchingSource(kind, sources))), nil
}

func (m *MockArtworkQueueRepo) SourcesInUse(kind model.Kind) ([]string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return nil, m.Err
	}
	sources := slice.Map(m.matchingSource(kind, nil), func(ia model.ItemArtwork) string { return ia.Source })
	return slice.Unique(sources), nil
}

func (m *MockArtworkQueueRepo) EnqueueBySource(kind model.Kind, sources []string, priority int) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return 0, m.Err
	}
	now := time.Now()
	var inserted int64
	for _, ia := range m.matchingSource(kind, sources) {
		k := iaKey(ia.ItemKind, ia.ItemID, ia.ImageType)
		if _, ok := m.Data[k]; ok { // DO NOTHING: never touch existing queue rows
			continue
		}
		m.Data[k] = model.ArtworkQueueItem{
			ItemKind:   ia.ItemKind,
			ItemID:     ia.ItemID,
			ImageType:  ia.ImageType,
			Priority:   priority,
			RetryAt:    now,
			EnqueuedAt: now,
		}
		inserted++
	}
	return inserted, nil
}

// EnqueueMissing mirrors the SQL set-difference insert: ExistingIDs[kind] minus ItemArtworkSource.
func (m *MockArtworkQueueRepo) EnqueueAllMissing(kind model.Kind, priority int) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.Err != nil {
		return 0, m.Err
	}
	hasRow := func(id string) bool {
		if m.ItemArtworkSource == nil {
			return false
		}
		for _, ia := range m.ItemArtworkSource.ItemData {
			if ia.ItemKind == kind.Prefix() && ia.ItemID == id {
				return true
			}
		}
		return false
	}
	now := time.Now()
	var inserted int64
	for id := range m.ExistingIDs[kind.Prefix()] {
		if hasRow(id) {
			continue
		}
		k := iaKey(kind.Prefix(), id, model.ImageTypePrimary)
		if _, ok := m.Data[k]; ok { // DO NOTHING: never touch existing queue rows
			continue
		}
		m.Data[k] = model.ArtworkQueueItem{
			ItemKind:   kind.Prefix(),
			ItemID:     id,
			ImageType:  model.ImageTypePrimary,
			Priority:   priority,
			RetryAt:    now,
			EnqueuedAt: now,
		}
		inserted++
	}
	return inserted, nil
}