go-agentic/allowance_error_test.go

package agentic

import (
	"errors"
	"testing"

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

// errorStore is a mock AllowanceStore that returns errors for specific operations.
type errorStore struct {
	*MemoryStore
	failIncrementUsage    bool
	failDecrementActive   bool
	failReturnTokens      bool
	failIncrementModel    bool
	failGetAllowance      bool
	failGetUsage          bool
}

func newErrorStore() *errorStore {
	return &errorStore{MemoryStore: NewMemoryStore()}
}

func (e *errorStore) GetAllowance(agentID string) (*AgentAllowance, error) {
	if e.failGetAllowance {
		return nil, errors.New("simulated GetAllowance error")
	}
	return e.MemoryStore.GetAllowance(agentID)
}

func (e *errorStore) GetUsage(agentID string) (*UsageRecord, error) {
	if e.failGetUsage {
		return nil, errors.New("simulated GetUsage error")
	}
	return e.MemoryStore.GetUsage(agentID)
}

func (e *errorStore) IncrementUsage(agentID string, tokens int64, jobs int) error {
	if e.failIncrementUsage {
		return errors.New("simulated IncrementUsage error")
	}
	return e.MemoryStore.IncrementUsage(agentID, tokens, jobs)
}

func (e *errorStore) DecrementActiveJobs(agentID string) error {
	if e.failDecrementActive {
		return errors.New("simulated DecrementActiveJobs error")
	}
	return e.MemoryStore.DecrementActiveJobs(agentID)
}

func (e *errorStore) ReturnTokens(agentID string, tokens int64) error {
	if e.failReturnTokens {
		return errors.New("simulated ReturnTokens error")
	}
	return e.MemoryStore.ReturnTokens(agentID, tokens)
}

func (e *errorStore) IncrementModelUsage(model string, tokens int64) error {
	if e.failIncrementModel {
		return errors.New("simulated IncrementModelUsage error")
	}
	return e.MemoryStore.IncrementModelUsage(model, tokens)
}

// --- RecordUsage error path tests ---

func TestRecordUsage_Bad_JobStarted_IncrementFails(t *testing.T) {
	store := newErrorStore()
	store.failIncrementUsage = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID: "agent-1",
		Event:   QuotaEventJobStarted,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to increment job count")
}

func TestRecordUsage_Bad_JobCompleted_IncrementFails(t *testing.T) {
	store := newErrorStore()
	store.failIncrementUsage = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		TokensIn:  100,
		TokensOut: 50,
		Event:     QuotaEventJobCompleted,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to record token usage")
}

func TestRecordUsage_Bad_JobCompleted_DecrementFails(t *testing.T) {
	store := newErrorStore()
	store.failDecrementActive = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		TokensIn:  100,
		TokensOut: 50,
		Event:     QuotaEventJobCompleted,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to decrement active jobs")
}

func TestRecordUsage_Bad_JobCompleted_ModelUsageFails(t *testing.T) {
	store := newErrorStore()
	store.failIncrementModel = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		Model:     "claude-sonnet",
		TokensIn:  100,
		TokensOut: 50,
		Event:     QuotaEventJobCompleted,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to record model usage")
}

func TestRecordUsage_Bad_JobFailed_IncrementFails(t *testing.T) {
	store := newErrorStore()
	store.failIncrementUsage = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		TokensIn:  100,
		TokensOut: 100,
		Event:     QuotaEventJobFailed,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to record token usage")
}

func TestRecordUsage_Bad_JobFailed_DecrementFails(t *testing.T) {
	store := newErrorStore()
	store.failDecrementActive = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		TokensIn:  100,
		TokensOut: 100,
		Event:     QuotaEventJobFailed,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to decrement active jobs")
}

func TestRecordUsage_Bad_JobFailed_ReturnTokensFails(t *testing.T) {
	store := newErrorStore()
	store.failReturnTokens = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		TokensIn:  100,
		TokensOut: 100,
		Event:     QuotaEventJobFailed,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to return tokens")
}

func TestRecordUsage_Bad_JobFailed_ModelUsageFails(t *testing.T) {
	store := newErrorStore()
	store.failIncrementModel = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		Model:     "claude-sonnet",
		TokensIn:  100,
		TokensOut: 100,
		Event:     QuotaEventJobFailed,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to record model usage")
}

func TestRecordUsage_Bad_JobCancelled_DecrementFails(t *testing.T) {
	store := newErrorStore()
	store.failDecrementActive = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		TokensIn:  100,
		TokensOut: 100,
		Event:     QuotaEventJobCancelled,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to decrement active jobs")
}

func TestRecordUsage_Bad_JobCancelled_ReturnTokensFails(t *testing.T) {
	store := newErrorStore()
	store.failReturnTokens = true
	svc := NewAllowanceService(store)

	err := svc.RecordUsage(UsageReport{
		AgentID:   "agent-1",
		TokensIn:  100,
		TokensOut: 100,
		Event:     QuotaEventJobCancelled,
	})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to return tokens")
}

// --- Check error path tests ---

func TestCheck_Bad_GetAllowanceFails(t *testing.T) {
	store := newErrorStore()
	store.failGetAllowance = true
	svc := NewAllowanceService(store)

	_, err := svc.Check("agent-1", "")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to get allowance")
}

func TestCheck_Bad_GetUsageFails(t *testing.T) {
	store := newErrorStore()
	svc := NewAllowanceService(store)

	_ = store.SetAllowance(&AgentAllowance{
		AgentID: "agent-1",
	})
	store.failGetUsage = true

	_, err := svc.Check("agent-1", "")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to get usage")
}

// --- ResetAgent error path ---

func TestResetAgent_Bad_ResetFails(t *testing.T) {
	// MemoryStore.ResetUsage never fails, but we can test the service
	// layer still returns nil for the happy path (already tested).
	// For a true error test, we'd need a mock, but the MemoryStore
	// never errors on ResetUsage. This confirms the pattern.
	store := NewMemoryStore()
	svc := NewAllowanceService(store)

	err := svc.ResetAgent("nonexistent-agent")
	require.NoError(t, err, "resetting a nonexistent agent should succeed")
}

// --- RecordUsage with unknown event type ---

func TestRecordUsage_Good_UnknownEvent(t *testing.T) {
	store := NewMemoryStore()
	svc := NewAllowanceService(store)

	// Unknown event should be a no-op (falls through the switch).
	err := svc.RecordUsage(UsageReport{
		AgentID: "agent-1",
		Event:   QuotaEvent("unknown_event"),
	})
	require.NoError(t, err, "unknown event should not error")
}