poimen-workflows/internal/batching/llm_batch_test.go

package batching

import (
	"testing"
	"time"

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

func TestLLMNewBatcher(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)
	assert.NotNil(t, batcher)
	assert.Equal(t, 0, batcher.QueueSize())
}

func TestLLMEnqueueRequest(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{
		ID:     "req-1",
		Type:   "implementer",
		Model:  "claude-opus",
		Prompt: "Generate code",
	}

	batcher.Enqueue(req)
	assert.Equal(t, 1, batcher.QueueSize())
}

func TestLLMEnqueueMultipleRequests(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	for i := 0; i < 5; i++ {
		req := &LLMRequest{
			Type:   "implementer",
			Model:  "claude-opus",
			Prompt: "Prompt",
		}
		batcher.Enqueue(req)
	}

	assert.Equal(t, 5, batcher.QueueSize())
}

func TestLLMAutoFlushOnMaxBatchSize(t *testing.T) {
	batcher := NewLLMBatcher(5, 10*time.Second)

	for i := 0; i < 5; i++ {
		req := &LLMRequest{
			Type:   "implementer",
			Model:  "claude-opus",
			Prompt: "Prompt",
		}
		batcher.Enqueue(req)
	}

	assert.Equal(t, 0, batcher.QueueSize())
	assert.Equal(t, 1, batcher.PendingBatchCount())
}

func TestLLMManualFlush(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{
		Type:   "implementer",
		Model:  "claude-opus",
		Prompt: "Prompt",
	}
	batcher.Enqueue(req)

	batcher.Flush()
	assert.Equal(t, 0, batcher.QueueSize())
	assert.Equal(t, 1, batcher.PendingBatchCount())
}

func TestLLMGetPendingBatch(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{
		Type:   "implementer",
		Model:  "claude-opus",
		Prompt: "Prompt",
	}
	batcher.Enqueue(req)
	batcher.Flush()

	batch := batcher.GetPendingBatch()
	assert.NotNil(t, batch)
	assert.Equal(t, 1, len(batch.Requests))
	assert.Equal(t, "implementer", batch.Type)
}

func TestLLMMarkBatchExecuting(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{Type: "implementer", Model: "claude-opus", Prompt: "test"}
	batcher.Enqueue(req)
	batcher.Flush()

	batch := batcher.GetPendingBatch()
	batcher.MarkBatchExecuting(batch.ID)

	updated := batcher.GetBatchByID(batch.ID)
	assert.Equal(t, "executing", updated.Status)
}

func TestLLMMarkBatchCompleted(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{
		ID:     "req-1",
		Type:   "implementer",
		Model:  "claude-opus",
		Prompt: "test",
	}
	batcher.Enqueue(req)
	batcher.Flush()

	batch := batcher.GetPendingBatch()

	results := map[string]*LLMResult{
		"req-1": {
			RequestID:  "req-1",
			Response:   "Generated code",
			TokenCount: 100,
		},
	}

	batcher.MarkBatchCompleted(batch.ID, results)

	executed := batcher.GetExecutedBatches()
	assert.Equal(t, 1, len(executed))
	assert.Equal(t, "completed", executed[0].Status)
}

func TestLLMMarkBatchFailed(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{
		ID:    "req-1",
		Type:  "implementer",
		Model: "claude-opus",
	}
	batcher.Enqueue(req)
	batcher.Flush()

	batch := batcher.GetPendingBatch()
	testErr := assert.AnError
	batcher.MarkBatchFailed(batch.ID, testErr)

	failed := batcher.GetBatchByID(batch.ID)
	assert.Equal(t, "failed", failed.Status)
	assert.Error(t, failed.Error)
}

func TestLLMGroupByTypeAndModel(t *testing.T) {
	batcher := NewLLMBatcher(100, 5*time.Second)

	// Add requests of different types
	for i := 0; i < 3; i++ {
		req := &LLMRequest{
			Type:   "implementer",
			Model:  "claude-opus",
			Prompt: "test",
		}
		batcher.Enqueue(req)
	}

	for i := 0; i < 2; i++ {
		req := &LLMRequest{
			Type:   "judge",
			Model:  "claude-opus",
			Prompt: "test",
		}
		batcher.Enqueue(req)
	}

	batcher.Flush()

	// Should create 2 batches (one for implementer, one for judge)
	assert.Equal(t, 2, batcher.PendingBatchCount())
}

func TestLLMGetStats(t *testing.T) {
	batcher := NewLLMBatcher(5, 5*time.Second)

	// Add 10 requests (will create 2 batches)
	for i := 0; i < 10; i++ {
		req := &LLMRequest{
			Type:   "implementer",
			Model:  "claude-opus",
			Prompt: "test",
		}
		batcher.Enqueue(req)
	}

	stats := batcher.GetStats()
	assert.Equal(t, 10, stats.TotalRequests)
	assert.Equal(t, 2, stats.TotalBatches)
	assert.Equal(t, 5.0, stats.AvgRequestsPerBatch)
	// 10 requests in 2 batches saves 8 API calls
	assert.Equal(t, 8, stats.APICallsSaved)
}

func TestLLMResultDelivery(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{
		ID:        "req-1",
		Type:      "implementer",
		Model:     "claude-opus",
		Prompt:    "test",
		ResultCh:  make(chan *LLMResult, 1),
	}

	batcher.Enqueue(req)
	batcher.Flush()

	batch := batcher.GetPendingBatch()

	results := map[string]*LLMResult{
		"req-1": {
			RequestID:  "req-1",
			Response:   "Response",
			TokenCount: 50,
		},
	}

	batcher.MarkBatchCompleted(batch.ID, results)

	// Check if result was delivered to channel
	select {
	case result := <-req.ResultCh:
		assert.NotNil(t, result)
		assert.Equal(t, "Response", result.Response)
	case <-time.After(1 * time.Second):
		t.Fatal("Result not delivered")
	}
}

func TestLLMMultipleBatches(t *testing.T) {
	batcher := NewLLMBatcher(3, 5*time.Second)

	// Create 3 batches (3 requests each)
	for batch := 0; batch < 3; batch++ {
		for i := 0; i < 3; i++ {
			req := &LLMRequest{
				Type:   "implementer",
				Model:  "claude-opus",
				Prompt: "test",
			}
			batcher.Enqueue(req)
		}
	}

	assert.Equal(t, 3, batcher.PendingBatchCount())
}

func TestLLMQueueSize(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
	batcher.Enqueue(req)

	assert.Equal(t, 1, batcher.QueueSize())
}

func TestLLMPendingBatchCount(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
	batcher.Enqueue(req)
	batcher.Flush()

	assert.Equal(t, 1, batcher.PendingBatchCount())
}

func TestLLMGetExecutedBatches(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	for i := 0; i < 2; i++ {
		req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
		batcher.Enqueue(req)
		batcher.Flush()

		batch := batcher.GetPendingBatch()
		batcher.MarkBatchCompleted(batch.ID, make(map[string]*LLMResult))
	}

	executed := batcher.GetExecutedBatches()
	assert.Equal(t, 2, len(executed))
}

func TestLLMTimeSinceLastFlush(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
	batcher.Enqueue(req)
	batcher.Flush()

	time.Sleep(100 * time.Millisecond)
	elapsed := batcher.TimeSinceLastFlush()

	assert.Greater(t, elapsed, 50*time.Millisecond)
	assert.Less(t, elapsed, 200*time.Millisecond)
}

func TestLLMShouldFlush(t *testing.T) {
	batcher := NewLLMBatcher(100, 100*time.Millisecond)

	req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
	batcher.Enqueue(req)

	// Should not flush yet
	assert.False(t, batcher.ShouldFlush())

	// Wait for age to exceed
	time.Sleep(150 * time.Millisecond)
	assert.True(t, batcher.ShouldFlush())
}

func TestLLMClear(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
	batcher.Enqueue(req)
	batcher.Flush()

	batcher.Clear()
	assert.Equal(t, 0, batcher.QueueSize())
	assert.Equal(t, 0, batcher.PendingBatchCount())
}

func TestLLMGetBatchByID(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
	batcher.Enqueue(req)
	batcher.Flush()

	batch := batcher.GetPendingBatch()
	retrieved := batcher.GetBatchByID(batch.ID)

	assert.NotNil(t, retrieved)
	assert.Equal(t, batch.ID, retrieved.ID)
}

func TestLLMGetBatchInfo(t *testing.T) {
	batch := &LLMBatch{
		ID:        "batch-1",
		Type:      "implementer",
		Model:     "claude-opus",
		Status:    "completed",
		CreatedAt: time.Now(),
	}

	info := batch.GetInfo()
	assert.Equal(t, "batch-1", info["id"])
	assert.Equal(t, "implementer", info["type"])
	assert.Equal(t, "claude-opus", info["model"])
	assert.Equal(t, "completed", info["status"])
}

func TestLLMEnqueueNil(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	batcher.Enqueue(nil)
	assert.Equal(t, 0, batcher.QueueSize())
}

func TestLLMAutoIDGeneration(t *testing.T) {
	req := &LLMRequest{
		Type:   "implementer",
		Model:  "claude-opus",
		Prompt: "test",
	}

	batcher := NewLLMBatcher(10, 5*time.Second)
	batcher.Enqueue(req)

	assert.NotEmpty(t, req.ID)
}

func TestLLMTokenCounting(t *testing.T) {
	batcher := NewLLMBatcher(10, 5*time.Second)

	req := &LLMRequest{
		ID:     "req-1",
		Type:   "implementer",
		Model:  "claude-opus",
		Prompt: "test",
	}
	batcher.Enqueue(req)
	batcher.Flush()

	batch := batcher.GetPendingBatch()

	results := map[string]*LLMResult{
		"req-1": {
			RequestID:  "req-1",
			Response:   "Response",
			TokenCount: 500,
		},
	}

	batcher.MarkBatchCompleted(batch.ID, results)

	stats := batcher.GetStats()
	assert.Equal(t, 500, stats.TotalTokens)
}

func BenchmarkLLMEnqueue(b *testing.B) {
	batcher := NewLLMBatcher(1000, 10*time.Second)

	for i := 0; i < b.N; i++ {
		req := &LLMRequest{
			Type:   "implementer",
			Model:  "claude-opus",
			Prompt: "test",
		}
		batcher.Enqueue(req)
	}
}

func BenchmarkLLMFlush(b *testing.B) {
	batcher := NewLLMBatcher(1000, 10*time.Second)

	for i := 0; i < b.N; i++ {
		req := &LLMRequest{Type: "implementer", Model: "claude-opus"}
		batcher.Enqueue(req)

		if (i + 1) % 100 == 0 {
			batcher.Flush()
		}
	}
}