owpengram-server/internal/store/postgres/message_store_testkit_test.go

package postgres

import (
	"context"
	"sync"
	"telesrv/internal/domain"
	"telesrv/internal/store/postgres/sqlcgen"
)

type testAllocatorBundle struct {
	boxIDs            *testBoxIDAllocator
	channelIDs        *testChannelIDAllocator
	channelMessageIDs *testChannelMessageIDAllocator
}

var testAllocatorBundles sync.Map

func testAllocatorsFor(db sqlcgen.DBTX) *testAllocatorBundle {
	if existing, ok := testAllocatorBundles.Load(db); ok {
		return existing.(*testAllocatorBundle)
	}
	bundle := &testAllocatorBundle{
		boxIDs:            &testBoxIDAllocator{source: NewMessageBoxCounterSource(db)},
		channelIDs:        &testChannelIDAllocator{source: NewChannelIDCounterSource(db)},
		channelMessageIDs: &testChannelMessageIDAllocator{source: NewChannelMessageIDCounterSource(db)},
	}
	actual, _ := testAllocatorBundles.LoadOrStore(db, bundle)
	return actual.(*testAllocatorBundle)
}

func newTestMessageStore(db sqlcgen.DBTX, opts ...MessageStoreOption) *MessageStore {
	bundle := testAllocatorsFor(db)
	all := append([]MessageStoreOption{WithMessageAllocators(bundle.boxIDs)}, opts...)
	return NewMessageStore(db, all...)
}

func newTestChannelStore(db sqlcgen.DBTX, opts ...ChannelStoreOption) *ChannelStore {
	bundle := testAllocatorsFor(db)
	all := append([]ChannelStoreOption{WithChannelAllocators(bundle.channelIDs, bundle.channelMessageIDs)}, opts...)
	return NewChannelStore(db, all...)
}

type testBoxIDAllocator struct {
	mu     sync.Mutex
	source *MessageBoxCounterSource
	values map[int64]int
}

func (a *testBoxIDAllocator) NextBoxID(ctx context.Context, userID int64) (int, error) {
	values, err := a.NextBoxIDs(ctx, []int64{userID})
	return values[userID], err
}

func (a *testBoxIDAllocator) NextBoxIDs(ctx context.Context, userIDs []int64) (map[int64]int, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	current, err := a.source.CurrentBatch(ctx, userIDs)
	if err != nil {
		return nil, err
	}
	if a.values == nil {
		a.values = make(map[int64]int, len(userIDs))
	}
	out := make(map[int64]int, len(userIDs))
	for _, userID := range userIDs {
		if current[userID] > a.values[userID] {
			a.values[userID] = current[userID]
		}
		a.values[userID]++
		out[userID] = a.values[userID]
	}
	return out, nil
}

func (a *testBoxIDAllocator) CurrentBoxID(ctx context.Context, userID int64) (int, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	current, err := a.source.Current(ctx, userID)
	if err != nil {
		return 0, err
	}
	if a.values[userID] > current {
		return a.values[userID], nil
	}
	return current, nil
}

type testChannelIDAllocator struct {
	mu      sync.Mutex
	source  *ChannelIDCounterSource
	current int64
}

func (a *testChannelIDAllocator) NextChannelID(ctx context.Context) (int64, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	current, err := a.source.Current(ctx, 0)
	if err != nil {
		return 0, err
	}
	if int64(current) > a.current {
		a.current = int64(current)
	}
	a.current++
	return a.current, nil
}

func (a *testChannelIDAllocator) CurrentChannelID(ctx context.Context) (int64, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	current, err := a.source.Current(ctx, 0)
	if err != nil {
		return 0, err
	}
	if int64(current) > a.current {
		a.current = int64(current)
	}
	return a.current, nil
}

type testChannelMessageIDAllocator struct {
	mu     sync.Mutex
	source *ChannelMessageIDCounterSource
	values map[int64]int
}

func (a *testChannelMessageIDAllocator) NextChannelMessageID(ctx context.Context, channelID int64) (int, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	current, err := a.source.Current(ctx, channelID)
	if err != nil {
		return 0, err
	}
	if a.values == nil {
		a.values = make(map[int64]int)
	}
	if current > a.values[channelID] {
		a.values[channelID] = current
	}
	a.values[channelID]++
	return a.values[channelID], nil
}

func (a *testChannelMessageIDAllocator) CurrentChannelMessageID(ctx context.Context, channelID int64) (int, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	current, err := a.source.Current(ctx, channelID)
	if err != nil {
		return 0, err
	}
	if a.values[channelID] > current {
		return a.values[channelID], nil
	}
	return current, nil
}

type fixedBoxIDAllocator struct {
	next int
}

func (a fixedBoxIDAllocator) NextBoxID(context.Context, int64) (int, error) {
	return a.next, nil
}

func (a fixedBoxIDAllocator) NextBoxIDs(_ context.Context, userIDs []int64) (map[int64]int, error) {
	out := make(map[int64]int, len(userIDs))
	for _, userID := range userIDs {
		out[userID] = a.next
	}
	return out, nil
}

func (a fixedBoxIDAllocator) CurrentBoxID(context.Context, int64) (int, error) {
	return a.next, nil
}

type perUserCounterAllocator struct {
	mu     sync.Mutex
	values map[int64]int
}

func (a *perUserCounterAllocator) NextBoxID(_ context.Context, userID int64) (int, error) {
	return a.next(userID), nil
}

func (a *perUserCounterAllocator) NextBoxIDs(_ context.Context, userIDs []int64) (map[int64]int, error) {
	out := make(map[int64]int, len(userIDs))
	for _, userID := range userIDs {
		if _, ok := out[userID]; !ok {
			out[userID] = a.next(userID)
		}
	}
	return out, nil
}

func (a *perUserCounterAllocator) CurrentBoxID(_ context.Context, userID int64) (int, error) {
	return a.current(userID), nil
}

func (a *perUserCounterAllocator) next(userID int64) int {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.values == nil {
		a.values = map[int64]int{}
	}
	a.values[userID]++
	return a.values[userID]
}

func (a *perUserCounterAllocator) current(userID int64) int {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.values[userID]
}

func messageIDs(messages []domain.Message) []int {
	out := make([]int, 0, len(messages))
	for _, msg := range messages {
		out = append(out, msg.ID)
	}
	return out
}

func sameInts(got, want []int) bool {
	if len(got) != len(want) {
		return false
	}
	for i := range got {
		if got[i] != want[i] {
			return false
		}
	}
	return true
}