go-agentic/allowance_sqlite_test.go

package agentic

import (
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// newTestSQLiteStore creates a SQLiteStore in a temp directory.
func newTestSQLiteStore(t *testing.T) *SQLiteStore {
	t.Helper()
	dbPath := filepath.Join(t.TempDir(), "test.db")
	s, err := NewSQLiteStore(dbPath)
	require.NoError(t, err)
	t.Cleanup(func() { _ = s.Close() })
	return s
}

// --- SetAllowance / GetAllowance ---

func TestSQLiteStore_SetGetAllowance_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	a := &AgentAllowance{
		AgentID:         "agent-1",
		DailyTokenLimit: 100000,
		DailyJobLimit:   10,
		ConcurrentJobs:  2,
		MaxJobDuration:  30 * time.Minute,
		ModelAllowlist:  []string{"claude-sonnet-4-5-20250929"},
	}

	err := s.SetAllowance(a)
	require.NoError(t, err)

	got, err := s.GetAllowance("agent-1")
	require.NoError(t, err)
	assert.Equal(t, a.AgentID, got.AgentID)
	assert.Equal(t, a.DailyTokenLimit, got.DailyTokenLimit)
	assert.Equal(t, a.DailyJobLimit, got.DailyJobLimit)
	assert.Equal(t, a.ConcurrentJobs, got.ConcurrentJobs)
	assert.Equal(t, a.MaxJobDuration, got.MaxJobDuration)
	assert.Equal(t, a.ModelAllowlist, got.ModelAllowlist)
}

func TestSQLiteStore_GetAllowance_Bad_NotFound(t *testing.T) {
	s := newTestSQLiteStore(t)
	_, err := s.GetAllowance("nonexistent")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "allowance not found")
}

func TestSQLiteStore_SetAllowance_Good_Overwrite(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.SetAllowance(&AgentAllowance{AgentID: "agent-1", DailyTokenLimit: 100})
	_ = s.SetAllowance(&AgentAllowance{AgentID: "agent-1", DailyTokenLimit: 200})

	got, err := s.GetAllowance("agent-1")
	require.NoError(t, err)
	assert.Equal(t, int64(200), got.DailyTokenLimit)
}

// --- GetUsage / IncrementUsage ---

func TestSQLiteStore_GetUsage_Good_Default(t *testing.T) {
	s := newTestSQLiteStore(t)

	u, err := s.GetUsage("agent-1")
	require.NoError(t, err)
	assert.Equal(t, "agent-1", u.AgentID)
	assert.Equal(t, int64(0), u.TokensUsed)
	assert.Equal(t, 0, u.JobsStarted)
	assert.Equal(t, 0, u.ActiveJobs)
}

func TestSQLiteStore_IncrementUsage_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	err := s.IncrementUsage("agent-1", 5000, 1)
	require.NoError(t, err)

	u, err := s.GetUsage("agent-1")
	require.NoError(t, err)
	assert.Equal(t, int64(5000), u.TokensUsed)
	assert.Equal(t, 1, u.JobsStarted)
	assert.Equal(t, 1, u.ActiveJobs)
}

func TestSQLiteStore_IncrementUsage_Good_Accumulates(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.IncrementUsage("agent-1", 1000, 1)
	_ = s.IncrementUsage("agent-1", 2000, 1)
	_ = s.IncrementUsage("agent-1", 3000, 0)

	u, err := s.GetUsage("agent-1")
	require.NoError(t, err)
	assert.Equal(t, int64(6000), u.TokensUsed)
	assert.Equal(t, 2, u.JobsStarted)
	assert.Equal(t, 2, u.ActiveJobs)
}

// --- DecrementActiveJobs ---

func TestSQLiteStore_DecrementActiveJobs_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.IncrementUsage("agent-1", 0, 2)
	_ = s.DecrementActiveJobs("agent-1")

	u, _ := s.GetUsage("agent-1")
	assert.Equal(t, 1, u.ActiveJobs)
}

func TestSQLiteStore_DecrementActiveJobs_Good_FloorAtZero(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.DecrementActiveJobs("agent-1") // no usage record yet
	_ = s.IncrementUsage("agent-1", 0, 0)
	_ = s.DecrementActiveJobs("agent-1") // should stay at 0

	u, _ := s.GetUsage("agent-1")
	assert.Equal(t, 0, u.ActiveJobs)
}

// --- ReturnTokens ---

func TestSQLiteStore_ReturnTokens_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.IncrementUsage("agent-1", 10000, 0)
	err := s.ReturnTokens("agent-1", 5000)
	require.NoError(t, err)

	u, _ := s.GetUsage("agent-1")
	assert.Equal(t, int64(5000), u.TokensUsed)
}

func TestSQLiteStore_ReturnTokens_Good_FloorAtZero(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.IncrementUsage("agent-1", 1000, 0)
	_ = s.ReturnTokens("agent-1", 5000) // more than used

	u, _ := s.GetUsage("agent-1")
	assert.Equal(t, int64(0), u.TokensUsed)
}

func TestSQLiteStore_ReturnTokens_Good_NoRecord(t *testing.T) {
	s := newTestSQLiteStore(t)

	// Return tokens for agent with no usage record -- should create one at 0
	err := s.ReturnTokens("agent-1", 500)
	require.NoError(t, err)

	u, _ := s.GetUsage("agent-1")
	assert.Equal(t, int64(0), u.TokensUsed)
}

// --- ResetUsage ---

func TestSQLiteStore_ResetUsage_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.IncrementUsage("agent-1", 50000, 5)

	err := s.ResetUsage("agent-1")
	require.NoError(t, err)

	u, _ := s.GetUsage("agent-1")
	assert.Equal(t, int64(0), u.TokensUsed)
	assert.Equal(t, 0, u.JobsStarted)
	assert.Equal(t, 0, u.ActiveJobs)
}

// --- ModelQuota ---

func TestSQLiteStore_GetModelQuota_Bad_NotFound(t *testing.T) {
	s := newTestSQLiteStore(t)
	_, err := s.GetModelQuota("nonexistent")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "model quota not found")
}

func TestSQLiteStore_SetGetModelQuota_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	q := &ModelQuota{
		Model:            "claude-opus-4-6",
		DailyTokenBudget: 500000,
		HourlyRateLimit:  100,
		CostCeiling:      10000,
	}
	err := s.SetModelQuota(q)
	require.NoError(t, err)

	got, err := s.GetModelQuota("claude-opus-4-6")
	require.NoError(t, err)
	assert.Equal(t, q.Model, got.Model)
	assert.Equal(t, q.DailyTokenBudget, got.DailyTokenBudget)
	assert.Equal(t, q.HourlyRateLimit, got.HourlyRateLimit)
	assert.Equal(t, q.CostCeiling, got.CostCeiling)
}

// --- ModelUsage ---

func TestSQLiteStore_ModelUsage_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.IncrementModelUsage("claude-sonnet", 10000)
	_ = s.IncrementModelUsage("claude-sonnet", 5000)

	usage, err := s.GetModelUsage("claude-sonnet")
	require.NoError(t, err)
	assert.Equal(t, int64(15000), usage)
}

func TestSQLiteStore_GetModelUsage_Good_Default(t *testing.T) {
	s := newTestSQLiteStore(t)

	usage, err := s.GetModelUsage("unknown-model")
	require.NoError(t, err)
	assert.Equal(t, int64(0), usage)
}

// --- Persistence: close and reopen ---

func TestSQLiteStore_Persistence_Good(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "persist.db")

	// Phase 1: write data
	s1, err := NewSQLiteStore(dbPath)
	require.NoError(t, err)

	_ = s1.SetAllowance(&AgentAllowance{
		AgentID:         "agent-1",
		DailyTokenLimit: 100000,
		MaxJobDuration:  15 * time.Minute,
	})
	_ = s1.IncrementUsage("agent-1", 25000, 3)
	_ = s1.SetModelQuota(&ModelQuota{Model: "claude-opus-4-6", DailyTokenBudget: 500000})
	_ = s1.IncrementModelUsage("claude-opus-4-6", 42000)
	require.NoError(t, s1.Close())

	// Phase 2: reopen and verify
	s2, err := NewSQLiteStore(dbPath)
	require.NoError(t, err)
	defer func() { _ = s2.Close() }()

	a, err := s2.GetAllowance("agent-1")
	require.NoError(t, err)
	assert.Equal(t, int64(100000), a.DailyTokenLimit)
	assert.Equal(t, 15*time.Minute, a.MaxJobDuration)

	u, err := s2.GetUsage("agent-1")
	require.NoError(t, err)
	assert.Equal(t, int64(25000), u.TokensUsed)
	assert.Equal(t, 3, u.JobsStarted)
	assert.Equal(t, 3, u.ActiveJobs)

	q, err := s2.GetModelQuota("claude-opus-4-6")
	require.NoError(t, err)
	assert.Equal(t, int64(500000), q.DailyTokenBudget)

	mu, err := s2.GetModelUsage("claude-opus-4-6")
	require.NoError(t, err)
	assert.Equal(t, int64(42000), mu)
}

// --- Concurrent access ---

func TestSQLiteStore_ConcurrentIncrementUsage_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	const goroutines = 10
	const tokensEach = 1000

	var wg sync.WaitGroup
	wg.Add(goroutines)
	for range goroutines {
		go func() {
			defer wg.Done()
			err := s.IncrementUsage("agent-1", tokensEach, 1)
			assert.NoError(t, err)
		}()
	}
	wg.Wait()

	u, err := s.GetUsage("agent-1")
	require.NoError(t, err)
	assert.Equal(t, int64(goroutines*tokensEach), u.TokensUsed)
	assert.Equal(t, goroutines, u.JobsStarted)
	assert.Equal(t, goroutines, u.ActiveJobs)
}

func TestSQLiteStore_ConcurrentModelUsage_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	const goroutines = 10
	const tokensEach int64 = 500

	var wg sync.WaitGroup
	wg.Add(goroutines)
	for range goroutines {
		go func() {
			defer wg.Done()
			err := s.IncrementModelUsage("claude-opus-4-6", tokensEach)
			assert.NoError(t, err)
		}()
	}
	wg.Wait()

	usage, err := s.GetModelUsage("claude-opus-4-6")
	require.NoError(t, err)
	assert.Equal(t, goroutines*tokensEach, usage)
}

func TestSQLiteStore_ConcurrentMixed_Good(t *testing.T) {
	s := newTestSQLiteStore(t)

	_ = s.SetAllowance(&AgentAllowance{
		AgentID:         "agent-1",
		DailyTokenLimit: 1000000,
		DailyJobLimit:   100,
		ConcurrentJobs:  50,
	})

	const goroutines = 10
	var wg sync.WaitGroup
	wg.Add(goroutines * 3)

	// Increment usage
	for range goroutines {
		go func() {
			defer wg.Done()
			_ = s.IncrementUsage("agent-1", 100, 1)
		}()
	}

	// Decrement active jobs
	for range goroutines {
		go func() {
			defer wg.Done()
			_ = s.DecrementActiveJobs("agent-1")
		}()
	}

	// Return tokens
	for range goroutines {
		go func() {
			defer wg.Done()
			_ = s.ReturnTokens("agent-1", 10)
		}()
	}

	wg.Wait()

	// Just verify no panics and data is consistent
	u, err := s.GetUsage("agent-1")
	require.NoError(t, err)
	assert.GreaterOrEqual(t, u.TokensUsed, int64(0))
	assert.GreaterOrEqual(t, u.ActiveJobs, 0)
}

// --- AllowanceService integration via SQLiteStore ---

func TestSQLiteStore_AllowanceServiceCheck_Good(t *testing.T) {
	s := newTestSQLiteStore(t)
	svc := NewAllowanceService(s)

	_ = s.SetAllowance(&AgentAllowance{
		AgentID:         "agent-1",
		DailyTokenLimit: 100000,
		DailyJobLimit:   10,
		ConcurrentJobs:  2,
	})

	result, err := svc.Check("agent-1", "")
	require.NoError(t, err)
	assert.True(t, result.Allowed)
	assert.Equal(t, AllowanceOK, result.Status)
}

func TestSQLiteStore_AllowanceServiceRecordUsage_Good(t *testing.T) {
	s := newTestSQLiteStore(t)
	svc := NewAllowanceService(s)

	_ = s.SetAllowance(&AgentAllowance{
		AgentID:         "agent-1",
		DailyTokenLimit: 100000,
	})

	// Start job
	err := svc.RecordUsage(UsageReport{
		AgentID: "agent-1",
		JobID:   "job-1",
		Event:   QuotaEventJobStarted,
	})
	require.NoError(t, err)

	// Complete job
	err = svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		JobID:     "job-1",
		Model:     "claude-sonnet",
		TokensIn:  1000,
		TokensOut: 500,
		Event:     QuotaEventJobCompleted,
	})
	require.NoError(t, err)

	u, _ := s.GetUsage("agent-1")
	assert.Equal(t, int64(1500), u.TokensUsed)
	assert.Equal(t, 0, u.ActiveJobs)
}

// --- Config-based factory ---

func TestNewAllowanceStoreFromConfig_Good_Memory(t *testing.T) {
	cfg := AllowanceConfig{StoreBackend: "memory"}
	s, err := NewAllowanceStoreFromConfig(cfg)
	require.NoError(t, err)
	_, ok := s.(*MemoryStore)
	assert.True(t, ok, "expected MemoryStore")
}

func TestNewAllowanceStoreFromConfig_Good_Default(t *testing.T) {
	cfg := AllowanceConfig{} // empty defaults to memory
	s, err := NewAllowanceStoreFromConfig(cfg)
	require.NoError(t, err)
	_, ok := s.(*MemoryStore)
	assert.True(t, ok, "expected MemoryStore for empty config")
}

func TestNewAllowanceStoreFromConfig_Good_SQLite(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "factory.db")
	cfg := AllowanceConfig{
		StoreBackend: "sqlite",
		StorePath:    dbPath,
	}
	s, err := NewAllowanceStoreFromConfig(cfg)
	require.NoError(t, err)
	ss, ok := s.(*SQLiteStore)
	assert.True(t, ok, "expected SQLiteStore")
	_ = ss.Close()
}

func TestNewAllowanceStoreFromConfig_Bad_UnknownBackend(t *testing.T) {
	cfg := AllowanceConfig{StoreBackend: "cassandra"}
	_, err := NewAllowanceStoreFromConfig(cfg)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "unsupported store backend")
}

// --- NewSQLiteStore error case ---

func TestNewSQLiteStore_Bad_InvalidPath(t *testing.T) {
	_, err := NewSQLiteStore("/nonexistent/deeply/nested/dir/test.db")
	require.Error(t, err)
}