poimen-workflows/internal/cost/cost_tracker_test.go

package cost

import (
	"testing"

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

func TestTrackLLMCost(t *testing.T) {
	tracker := NewCostTracker()
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	
	entries := tracker.GetEntries()
	assert.Equal(t, 1, len(entries))
	assert.Equal(t, "llm", entries[0].Type)
	assert.Equal(t, 0.1, entries[0].Amount) // 1000 tokens * 0.0001
}

func TestTrackGitCost(t *testing.T) {
	tracker := NewCostTracker()
	tracker.TrackGitCost("wf-1", 5)
	
	entries := tracker.GetEntries()
	assert.Equal(t, 1, len(entries))
	assert.Equal(t, "git", entries[0].Type)
}

func TestTrackComputeCost(t *testing.T) {
	tracker := NewCostTracker()
	tracker.TrackComputeCost("wf-1", "task-1", 3600000) // 1 hour in ms
	
	entries := tracker.GetEntries()
	assert.Equal(t, 1, len(entries))
	assert.Equal(t, "compute", entries[0].Type)
	assert.Equal(t, 0.5, entries[0].Amount) // 1 hour * $0.5/hour
}

func TestGetTotalCost(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	tracker.TrackComputeCost("wf-1", "task-1", 3600000)
	
	total := tracker.GetTotalCost()
	assert.Equal(t, 0.6, total) // 0.1 + 0.5
}

func TestGetWorkflowCost(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	tracker.TrackLLMCost("wf-2", "task-1", 2000)
	
	cost := tracker.GetWorkflowCost("wf-1")
	assert.Equal(t, 0.1, cost)
	
	cost = tracker.GetWorkflowCost("wf-2")
	assert.Equal(t, 0.2, cost)
}

func TestGetCostByType(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	tracker.TrackLLMCost("wf-1", "task-2", 1000)
	tracker.TrackComputeCost("wf-1", "task-3", 3600000)
	
	llmCost := tracker.GetCostByType("llm")
	assert.Equal(t, 0.2, llmCost)
	
	computeCost := tracker.GetCostByType("compute")
	assert.Equal(t, 0.5, computeCost)
}

func TestGetAverageCostPerTask(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	tracker.TrackLLMCost("wf-1", "task-2", 1000)
	
	avg := tracker.GetAverageCostPerTask("wf-1")
	assert.Equal(t, 0.1, avg)
}

func TestGetOptimizationSuggestions(t *testing.T) {
	tracker := NewCostTracker()
	
	// High LLM cost
	tracker.TrackLLMCost("wf-1", "task-1", 10000)
	tracker.TrackLLMCost("wf-1", "task-2", 10000)
	tracker.TrackComputeCost("wf-1", "task-3", 360000) // 0.1 seconds
	
	suggestions := tracker.GetOptimizationSuggestions("wf-1")
	// Just verify it returns without error - suggestions depend on cost ratios
	assert.NotNil(t, suggestions)
}

func TestGetEntriesForWorkflow(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	tracker.TrackLLMCost("wf-2", "task-1", 1000)
	
	entries := tracker.GetEntriesForWorkflow("wf-1")
	assert.Equal(t, 1, len(entries))
}

func TestSetAndGetRate(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.SetRate("custom", 0.5)
	rate := tracker.GetRate("custom")
	
	assert.Equal(t, 0.5, rate)
}

func TestClear(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	tracker.Clear()
	
	entries := tracker.GetEntries()
	assert.Equal(t, 0, len(entries))
}

func TestMultipleCosts(t *testing.T) {
	tracker := NewCostTracker()
	
	tracker.TrackLLMCost("wf-1", "task-1", 1000)
	tracker.TrackGitCost("wf-1", 5)
	tracker.TrackComputeCost("wf-1", "task-1", 1800000) // 30 min
	
	total := tracker.GetTotalCost()
	assert.True(t, total > 0.2)
}

func TestZeroCost(t *testing.T) {
	tracker := NewCostTracker()
	
	cost := tracker.GetWorkflowCost("nonexistent")
	assert.Equal(t, 0.0, cost)
}

func BenchmarkTrackLLMCost(b *testing.B) {
	tracker := NewCostTracker()
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tracker.TrackLLMCost("wf-1", "task-1", 1000)
	}
}