poimen-workflows/internal/routing/llm_router_test.go

package routing

import (
	"encoding/json"
	"testing"
)

func TestParseIntentResponse(t *testing.T) {
	tests := []struct {
		name     string
		response string
		wantErr  bool
		validate func(*testing.T, *Intent)
	}{
		{
			name: "basic intent",
			response: `{
				"activities": ["CloneRepoActivity", "AnalyzeCodeActivity"],
				"parameters": {"repo": "https://github.com/test/repo"},
				"isCron": false,
				"workflowName": "analyze-repo"
			}`,
			wantErr: false,
			validate: func(t *testing.T, intent *Intent) {
				if len(intent.Activities) != 2 {
					t.Errorf("expected 2 activities, got %d", len(intent.Activities))
				}
				if intent.Activities[0] != "CloneRepoActivity" {
					t.Errorf("expected CloneRepoActivity first, got %s", intent.Activities[0])
				}
				if intent.IsCron {
					t.Error("expected isCron=false")
				}
			},
		},
		{
			name: "cron intent",
			response: `{
				"activities": ["CloneRepoActivity", "SecurityScanActivity"],
				"parameters": {"repo": "https://github.com/test/repo"},
				"isCron": true,
				"cronSchedule": "0 2 * * *",
				"cronTimezone": "America/New_York",
				"workflowName": "daily-security-scan"
			}`,
			wantErr: false,
			validate: func(t *testing.T, intent *Intent) {
				if !intent.IsCron {
					t.Error("expected isCron=true")
				}
				if intent.CronSchedule != "0 2 * * *" {
					t.Errorf("expected cron schedule '0 2 * * *', got %s", intent.CronSchedule)
				}
				if intent.CronTimezone != "America/New_York" {
					t.Errorf("expected timezone 'America/New_York', got %s", intent.CronTimezone)
				}
			},
		},
		{
			name: "with markdown code block",
			response: "```json\n{\"activities\": [\"CloneRepoActivity\"], \"parameters\": {}, \"isCron\": false}\n```",
			wantErr: false,
			validate: func(t *testing.T, intent *Intent) {
				if len(intent.Activities) != 1 {
					t.Errorf("expected 1 activity, got %d", len(intent.Activities))
				}
			},
		},
		{
			name: "defaults applied",
			response: `{"activities": ["CloneRepoActivity"], "parameters": {}}`,
			wantErr: false,
			validate: func(t *testing.T, intent *Intent) {
				if intent.CronTimezone != "UTC" {
					t.Errorf("expected default timezone UTC, got %s", intent.CronTimezone)
				}
				if intent.ErrorHandling != "retry" {
					t.Errorf("expected default errorHandling 'retry', got %s", intent.ErrorHandling)
				}
				if intent.WorkflowName != "generated-workflow" {
					t.Errorf("expected default workflowName, got %s", intent.WorkflowName)
				}
			},
		},
		{
			name:    "invalid json",
			response: "this is not json",
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			intent, err := parseIntentResponse(tt.response)
			if tt.wantErr {
				if err == nil {
					t.Error("expected error, got nil")
				}
				return
			}
			if err != nil {
				t.Errorf("unexpected error: %v", err)
				return
			}
			if tt.validate != nil {
				tt.validate(t, intent)
			}
		})
	}
}

func TestBuildSpec(t *testing.T) {
	// Load knowledge base
	kb, err := LoadKnowledgeBaseFromDefaultPath()
	if err != nil {
		t.Fatalf("failed to load knowledge base: %v", err)
	}

	router := &LLMRouter{
		knowledgeBase: kb,
	}

	intent := &Intent{
		Activities:    []string{"CloneRepoActivity", "AnalyzeCodeActivity", "SecurityScanActivity"},
		Parameters:    map[string]interface{}{"repo": "https://github.com/test/repo", "branch": "main"},
		WorkflowName:  "test-workflow",
		ErrorHandling: "retry",
	}

	input := LLMRouterInput{
		Message: "Analyze repo for security",
		Context: map[string]interface{}{},
	}

	spec, err := router.buildSpec(intent, input)
	if err != nil {
		t.Fatalf("buildSpec failed: %v", err)
	}

	// Validate spec
	if spec.Name != "test-workflow" {
		t.Errorf("expected name 'test-workflow', got %s", spec.Name)
	}

	if len(spec.States) < 3 {
		t.Errorf("expected at least 3 states, got %d", len(spec.States))
	}

	// First state should be CloneRepoActivity
	if spec.States[0].Resource != "CloneRepoActivity" {
		t.Errorf("expected first state to be CloneRepoActivity, got %s", spec.States[0].Resource)
	}

	// Last activity state should have End=true
	lastActivityIdx := len(spec.States) - 1
	if spec.States[lastActivityIdx].Type == StateTypeFail {
		lastActivityIdx--
	}
	if !spec.States[lastActivityIdx].End {
		t.Error("expected last activity state to have End=true")
	}

	// Check retry policy on flaky activity (AnalyzeCodeActivity)
	for _, state := range spec.States {
		if state.Resource == "AnalyzeCodeActivity" {
			if state.Retry == nil {
				t.Error("expected retry policy on flaky activity")
			} else if state.Retry.MaxAttempts != 3 {
				t.Errorf("expected 3 max attempts for flaky activity, got %d", state.Retry.MaxAttempts)
			}
			if len(state.Catch) == 0 {
				t.Error("expected catch clause on flaky activity")
			}
		}
	}
}

func TestBuildCronSpec(t *testing.T) {
	kb, err := LoadKnowledgeBaseFromDefaultPath()
	if err != nil {
		t.Fatalf("failed to load knowledge base: %v", err)
	}

	router := &LLMRouter{
		knowledgeBase: kb,
	}

	intent := &Intent{
		Activities:   []string{"CloneRepoActivity", "SecurityScanActivity"},
		Parameters:   map[string]interface{}{"repo": "https://github.com/test/repo"},
		IsCron:       true,
		CronSchedule: "0 2 * * *",
		CronTimezone: "UTC",
		WorkflowName: "daily-scan",
	}

	input := LLMRouterInput{
		Message: "Run security scan daily at 2 AM",
	}

	cronSpec, err := router.buildCronSpec(intent, input)
	if err != nil {
		t.Fatalf("buildCronSpec failed: %v", err)
	}

	if cronSpec.Type != "CronWorkflow" {
		t.Errorf("expected type 'CronWorkflow', got %s", cronSpec.Type)
	}
	if cronSpec.Schedule != "0 2 * * *" {
		t.Errorf("expected schedule '0 2 * * *', got %s", cronSpec.Schedule)
	}
	if cronSpec.Timezone != "UTC" {
		t.Errorf("expected timezone 'UTC', got %s", cronSpec.Timezone)
	}
	if !cronSpec.EnableHistory {
		t.Error("expected EnableHistory=true")
	}
}

func TestBuildParameters(t *testing.T) {
	kb, err := LoadKnowledgeBaseFromDefaultPath()
	if err != nil {
		t.Fatalf("failed to load knowledge base: %v", err)
	}

	router := &LLMRouter{
		knowledgeBase: kb,
	}

	// Test first activity (CloneRepoActivity) - should use input references
	cloneAct := kb.GetActivity("CloneRepoActivity")
	intent := &Intent{
		Activities: []string{"CloneRepoActivity", "AnalyzeCodeActivity"},
		Parameters: map[string]interface{}{"repo": "https://github.com/test/repo"},
	}

	params := router.buildParameters(cloneAct, intent, 0)
	if params["repo"] != "https://github.com/test/repo" {
		t.Errorf("expected repo from parameters, got %v", params["repo"])
	}

	// Test second activity (AnalyzeCodeActivity) - should reference previous output
	analyzeAct := kb.GetActivity("AnalyzeCodeActivity")
	params = router.buildParameters(analyzeAct, intent, 1)
	if params["path"] != "${CloneRepoActivity.output.path}" {
		t.Errorf("expected JSONPath reference to CloneRepoActivity.output.path, got %v", params["path"])
	}
}

func TestBuildRetryPolicy(t *testing.T) {
	kb, err := LoadKnowledgeBaseFromDefaultPath()
	if err != nil {
		t.Fatalf("failed to load knowledge base: %v", err)
	}

	router := &LLMRouter{
		knowledgeBase: kb,
	}

	// Flaky activity with retry error handling
	analyzeAct := kb.GetActivity("AnalyzeCodeActivity")
	intent := &Intent{ErrorHandling: "retry"}
	policy := router.buildRetryPolicy(analyzeAct, intent)

	if policy.MaxAttempts != 3 {
		t.Errorf("expected 3 max attempts for flaky activity, got %d", policy.MaxAttempts)
	}
	if policy.BackoffRate != 2.0 {
		t.Errorf("expected backoff rate 2.0, got %f", policy.BackoffRate)
	}

	// Fail-fast error handling
	intent = &Intent{ErrorHandling: "fail-fast"}
	policy = router.buildRetryPolicy(analyzeAct, intent)

	if policy.MaxAttempts != 1 {
		t.Errorf("expected 1 max attempt for fail-fast, got %d", policy.MaxAttempts)
	}
}

func TestIntentJSONMarshal(t *testing.T) {
	intent := &Intent{
		Activities:    []string{"CloneRepoActivity"},
		Parameters:    map[string]interface{}{"repo": "https://test"},
		IsCron:        true,
		CronSchedule:  "0 * * * *",
		CronTimezone:  "UTC",
		WorkflowName:  "test",
		ErrorHandling: "retry",
	}

	data, err := json.Marshal(intent)
	if err != nil {
		t.Fatalf("marshal failed: %v", err)
	}

	var decoded Intent
	if err := json.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("unmarshal failed: %v", err)
	}

	if decoded.CronSchedule != intent.CronSchedule {
		t.Errorf("expected schedule %s, got %s", intent.CronSchedule, decoded.CronSchedule)
	}
}