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 TestGetStringFromMap(t *testing.T) {
	tests := []struct {
		name     string
		m        map[string]interface{}
		key      string
		expected string
	}{
		{"nil map", nil, "key", ""},
		{"missing key", map[string]interface{}{"a": "b"}, "key", ""},
		{"found string", map[string]interface{}{"key": "value"}, "key", "value"},
		{"non-string value", map[string]interface{}{"key": 123}, "key", ""},
		{"empty string", map[string]interface{}{"key": ""}, "key", ""},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := getStringFromMap(tt.m, tt.key)
			if got != tt.expected {
				t.Errorf("getStringFromMap() = %q, want %q", got, tt.expected)
			}
		})
	}
}

func TestFirstNonEmpty(t *testing.T) {
	tests := []struct {
		name     string
		values   []string
		expected string
	}{
		{"all empty", []string{"", "", ""}, ""},
		{"first non-empty", []string{"first", "second"}, "first"},
		{"second non-empty", []string{"", "second", "third"}, "second"},
		{"last non-empty", []string{"", "", "last"}, "last"},
		{"no values", []string{}, ""},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := firstNonEmpty(tt.values...)
			if got != tt.expected {
				t.Errorf("firstNonEmpty() = %q, want %q", got, tt.expected)
			}
		})
	}
}

func TestIsCommonInputField(t *testing.T) {
	tests := []struct {
		name     string
		expected bool
	}{
		{"repo", true},
		{"path", true},
		{"branch", true},
		{"unknown", false},
		{"Repository", false}, // case-sensitive
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := isCommonInputField(tt.name)
			if got != tt.expected {
				t.Errorf("isCommonInputField(%q) = %v, want %v", tt.name, got, tt.expected)
			}
		})
	}
}

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

	router := &LLMRouter{knowledgeBase: kb}

	tests := []struct {
		name           string
		intentSchedule string
		intentTimezone string
		paramSchedule  string
		paramTimezone  string
		wantSchedule   string
		wantTimezone   string
	}{
		{
			name:           "from intent",
			intentSchedule: "0 2 * * *",
			intentTimezone: "PST",
			wantSchedule:   "0 2 * * *",
			wantTimezone:   "PST",
		},
		{
			name:          "from params",
			paramSchedule: "0 3 * * *",
			paramTimezone: "EST",
			wantSchedule:  "0 3 * * *",
			wantTimezone:  "EST",
		},
		{
			name:         "default UTC",
			wantSchedule: "",
			wantTimezone: "UTC",
		},
		{
			name:           "intent priority",
			intentSchedule: "0 1 * * *",
			intentTimezone: "UTC",
			paramSchedule:  "0 2 * * *",
			paramTimezone:  "PST",
			wantSchedule:   "0 1 * * *",
			wantTimezone:   "UTC",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			intent := &Intent{
				Activities:   []string{"CloneRepoActivity"},
				Parameters:   map[string]interface{}{},
				IsCron:       true,
				CronSchedule: tt.intentSchedule,
				CronTimezone: tt.intentTimezone,
				WorkflowName: "test",
			}
			if tt.paramSchedule != "" {
				intent.Parameters["cronSchedule"] = tt.paramSchedule
			}
			if tt.paramTimezone != "" {
				intent.Parameters["cronTimezone"] = tt.paramTimezone
			}

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

			if spec.Schedule != tt.wantSchedule {
				t.Errorf("schedule = %q, want %q", spec.Schedule, tt.wantSchedule)
			}
			if spec.Timezone != tt.wantTimezone {
				t.Errorf("timezone = %q, want %q", spec.Timezone, tt.wantTimezone)
			}
		})
	}
}

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

	resolver := &paramResolver{
		intent: &Intent{
			Activities: []string{"CloneRepoActivity", "AnalyzeCodeActivity"},
			Parameters: map[string]interface{}{},
		},
		kb:        kb,
		prevState: "CloneRepoActivity",
	}

	// Should find matching output from CloneRepoActivity
	got := resolver.fromPrevOutput("path")
	if got == nil {
		t.Error("expected to find path from prev output")
	}
	if got != "${CloneRepoActivity.output.path}" {
		t.Errorf("got %v, want ${CloneRepoActivity.output.path}", got)
	}

	// Should not find non-existent output
	got = resolver.fromPrevOutput("nonexistent")
	if got != nil {
		t.Errorf("expected nil for nonexistent, got %v", got)
	}

	// Empty prevState should return nil
	resolver.prevState = ""
	got = resolver.fromPrevOutput("path")
	if got != nil {
		t.Errorf("expected nil for empty prevState, got %v", got)
	}
}

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)
	}
}