AIRouter/internal/task/task_test.go

package task

import (
	"testing"
	"time"
)

func TestNewTask(t *testing.T) {
	task := NewTask("task-1", "req-1", "app-1", "tenant-1", "general-chat", PriorityNormal, false)
	if task.ID != "task-1" {
		t.Errorf("expected ID task-1, got %s", task.ID)
	}
	if task.State != StateQueued {
		t.Errorf("expected state QUEUED, got %s", task.State)
	}
	if task.Priority != PriorityNormal {
		t.Errorf("expected priority P2, got %d", task.Priority)
	}
}

func TestValidTransitions(t *testing.T) {
	tests := []struct {
		from TaskState
		to   TaskState
		ok   bool
	}{
		{StateQueued, StateRunning, true},
		{StateQueued, StateFailed, true},
		{StateQueued, StateCancelled, true},
		{StateQueued, StateCompleted, false},
		{StateRunning, StateStreaming, true},
		{StateRunning, StateCompleted, true},
		{StateRunning, StateFailed, true},
		{StateRunning, StateCancelled, true},
		{StateRunning, StateQueued, false},
		{StateStreaming, StateCompleted, true},
		{StateStreaming, StateFailed, true},
		{StateStreaming, StateCancelled, true},
		{StateStreaming, StateRunning, false},
		{StateCompleted, StateRunning, false},
		{StateFailed, StateCompleted, false},
		{StateCancelled, StateRunning, false},
	}

	for _, tt := range tests {
		task := &Task{State: tt.from, cancelCh: make(chan struct{})}
		err := task.Transition(tt.to)
		if tt.ok && err != nil {
			t.Errorf("expected %s -> %s to succeed, got error: %v", tt.from, tt.to, err)
		}
		if !tt.ok && err == nil {
			t.Errorf("expected %s -> %s to fail, but it succeeded", tt.from, tt.to)
		}
	}
}

func TestTaskCancel(t *testing.T) {
	task := NewTask("task-1", "req-1", "app-1", "tenant-1", "general-chat", PriorityNormal, false)

	if task.IsCancelled() {
		t.Error("task should not be cancelled initially")
	}

	err := task.Cancel("client_disconnect")
	if err != nil {
		t.Errorf("cancel failed: %v", err)
	}

	if !task.IsCancelled() {
		t.Error("task should be cancelled after Cancel()")
	}

	if task.GetState() != StateCancelled {
		t.Errorf("expected state CANCELLED, got %s", task.GetState())
	}

	if task.CancelReason != "client_disconnect" {
		t.Errorf("expected cancel reason 'client_disconnect', got %s", task.CancelReason)
	}

	// Cancel again should fail
	err = task.Cancel("second_attempt")
	if err == nil {
		t.Error("expected error on double cancel")
	}
}

func TestTaskCancelledChannel(t *testing.T) {
	task := NewTask("task-1", "req-1", "app-1", "tenant-1", "general-chat", PriorityNormal, false)

	select {
	case <-task.Cancelled():
		t.Error("channel should not be closed before cancel")
	default:
	}

	task.Cancel("test")

	select {
	case <-task.Cancelled():
		// expected
	case <-time.After(100 * time.Millisecond):
		t.Error("channel should be closed after cancel")
	}
}

func TestIsTerminal(t *testing.T) {
	tests := []struct {
		state    TaskState
		terminal bool
	}{
		{StateQueued, false},
		{StateRunning, false},
		{StateStreaming, false},
		{StateCompleted, true},
		{StateFailed, true},
		{StateCancelled, true},
	}

	for _, tt := range tests {
		task := &Task{State: tt.state}
		if task.IsTerminal() != tt.terminal {
			t.Errorf("expected IsTerminal()=%v for state %s, got %v", tt.terminal, tt.state, task.IsTerminal())
		}
	}
}

func TestTransitionSetsTimestamps(t *testing.T) {
	task := &Task{State: StateQueued, cancelCh: make(chan struct{})}

	err := task.Transition(StateRunning)
	if err != nil {
		t.Fatalf("transition to RUNNING failed: %v", err)
	}
	if task.StartedAt == nil {
		t.Error("expected StartedAt to be set after transition to RUNNING")
	}

	err = task.Transition(StateCompleted)
	if err != nil {
		t.Fatalf("transition to COMPLETED failed: %v", err)
	}
	if task.CompletedAt == nil {
		t.Error("expected CompletedAt to be set after transition to COMPLETED")
	}
}