path: root/internal/store/agents_test.go

package store

import (
	"database/sql"
	"path/filepath"
	"testing"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"task-dashboard/internal/models"
)

// =============================================================================
// Agent Session Tests
// =============================================================================

func setupTestStoreWithAgents(t *testing.T) *Store {
	t.Helper()

	tempDir := t.TempDir()
	dbPath := filepath.Join(tempDir, "test.db")

	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		t.Fatalf("Failed to open test database: %v", err)
	}

	db.SetMaxOpenConns(1)
	store := &Store{db: db}

	schema := `
		CREATE TABLE IF NOT EXISTS agents (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			name TEXT NOT NULL,
			agent_id TEXT UNIQUE NOT NULL,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			last_seen DATETIME,
			trusted BOOLEAN DEFAULT 1
		);
		CREATE TABLE IF NOT EXISTS agent_sessions (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			request_token TEXT UNIQUE NOT NULL,
			agent_name TEXT NOT NULL,
			agent_id TEXT NOT NULL,
			status TEXT DEFAULT 'pending',
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			expires_at DATETIME NOT NULL,
			session_token TEXT,
			session_expires_at DATETIME
		);
	`
	if _, err := db.Exec(schema); err != nil {
		t.Fatalf("Failed to create schema: %v", err)
	}

	return store
}

func TestAgentSession_CreateAndRetrieve(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	expiresAt := time.Now().Add(5 * time.Minute)
	session := &models.AgentSession{
		RequestToken: "req-token-123",
		AgentName:    "TestAgent",
		AgentID:      "agent-uuid-123",
		ExpiresAt:    expiresAt,
	}

	// Create session
	if err := store.CreateAgentSession(session); err != nil {
		t.Fatalf("Failed to create session: %v", err)
	}
	if session.ID == 0 {
		t.Error("Session ID should be set after create")
	}

	// Get by request token
	retrieved, err := store.GetAgentSessionByRequestToken("req-token-123")
	if err != nil {
		t.Fatalf("Failed to get session: %v", err)
	}
	if retrieved == nil {
		t.Fatal("Expected session to exist")
	}
	if retrieved.AgentName != "TestAgent" {
		t.Errorf("Expected name 'TestAgent', got '%s'", retrieved.AgentName)
	}
	if retrieved.Status != "pending" {
		t.Errorf("Expected status 'pending', got '%s'", retrieved.Status)
	}

	// Get pending by agent ID
	pending, err := store.GetPendingAgentSessionByAgentID("agent-uuid-123")
	if err != nil {
		t.Fatalf("Failed to get pending session: %v", err)
	}
	if pending == nil {
		t.Fatal("Expected pending session to exist")
	}
}

func TestAgentSession_ApproveAndDeny(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	// Create two sessions
	session1 := &models.AgentSession{
		RequestToken: "approve-token",
		AgentName:    "Agent1",
		AgentID:      "agent-1",
		ExpiresAt:    time.Now().Add(5 * time.Minute),
	}
	session2 := &models.AgentSession{
		RequestToken: "deny-token",
		AgentName:    "Agent2",
		AgentID:      "agent-2",
		ExpiresAt:    time.Now().Add(5 * time.Minute),
	}
	_ = store.CreateAgentSession(session1)
	_ = store.CreateAgentSession(session2)

	// Approve session1
	sessionExpiry := time.Now().Add(1 * time.Hour)
	if err := store.ApproveAgentSession("approve-token", "session-token-abc", sessionExpiry); err != nil {
		t.Fatalf("Failed to approve session: %v", err)
	}

	// Verify approval
	approved, _ := store.GetAgentSessionByRequestToken("approve-token")
	if approved.Status != "approved" {
		t.Errorf("Expected status 'approved', got '%s'", approved.Status)
	}
	if approved.SessionToken != "session-token-abc" {
		t.Errorf("Expected session token 'session-token-abc', got '%s'", approved.SessionToken)
	}

	// Deny session2
	if err := store.DenyAgentSession("deny-token"); err != nil {
		t.Fatalf("Failed to deny session: %v", err)
	}

	denied, _ := store.GetAgentSessionByRequestToken("deny-token")
	if denied.Status != "denied" {
		t.Errorf("Expected status 'denied', got '%s'", denied.Status)
	}
}

func TestAgentSession_GetBySessionToken(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	session := &models.AgentSession{
		RequestToken: "req-for-session",
		AgentName:    "SessionAgent",
		AgentID:      "session-agent",
		ExpiresAt:    time.Now().Add(5 * time.Minute),
	}
	_ = store.CreateAgentSession(session)
	_ = store.ApproveAgentSession("req-for-session", "active-session", time.Now().Add(1*time.Hour))

	// Get by session token
	retrieved, err := store.GetAgentSessionBySessionToken("active-session")
	if err != nil {
		t.Fatalf("Failed to get by session token: %v", err)
	}
	if retrieved == nil {
		t.Fatal("Expected session to exist")
	}
	if retrieved.AgentName != "SessionAgent" {
		t.Errorf("Expected 'SessionAgent', got '%s'", retrieved.AgentName)
	}
}

func TestAgentSession_GetPending(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	// Create pending sessions
	for i := 0; i < 3; i++ {
		session := &models.AgentSession{
			RequestToken: "pending-" + string(rune('0'+i)),
			AgentName:    "Agent" + string(rune('0'+i)),
			AgentID:      "agent-" + string(rune('0'+i)),
			ExpiresAt:    time.Now().Add(5 * time.Minute),
		}
		_ = store.CreateAgentSession(session)
	}

	// Get pending sessions
	pending, err := store.GetPendingAgentSessions()
	if err != nil {
		t.Fatalf("Failed to get pending sessions: %v", err)
	}
	if len(pending) != 3 {
		t.Errorf("Expected 3 pending sessions, got %d", len(pending))
	}
}

func TestAgentSession_Invalidate(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	// Create sessions for same agent
	for i := 0; i < 2; i++ {
		session := &models.AgentSession{
			RequestToken: "inv-" + string(rune('0'+i)),
			AgentName:    "SameAgent",
			AgentID:      "same-agent",
			ExpiresAt:    time.Now().Add(5 * time.Minute),
		}
		_ = store.CreateAgentSession(session)
	}

	// Invalidate all sessions for agent
	if err := store.InvalidatePreviousAgentSessions("same-agent"); err != nil {
		t.Fatalf("Failed to invalidate sessions: %v", err)
	}

	// Verify no pending sessions
	pending, _ := store.GetPendingAgentSessions()
	for _, s := range pending {
		if s.AgentID == "same-agent" {
			t.Error("Session should be invalidated")
		}
	}
}

// =============================================================================
// Agent Tests
// =============================================================================

func TestAgent_CreateAndRetrieve(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	// Create agent
	if err := store.CreateOrUpdateAgent("TestBot", "bot-uuid-123"); err != nil {
		t.Fatalf("Failed to create agent: %v", err)
	}

	// Get by agent ID
	agent, err := store.GetAgentByAgentID("bot-uuid-123")
	if err != nil {
		t.Fatalf("Failed to get agent: %v", err)
	}
	if agent == nil {
		t.Fatal("Expected agent to exist")
	}
	if agent.Name != "TestBot" {
		t.Errorf("Expected name 'TestBot', got '%s'", agent.Name)
	}
	if !agent.Trusted {
		t.Error("New agent should be trusted by default")
	}

	// Get by name
	byName, err := store.GetAgentByName("TestBot")
	if err != nil {
		t.Fatalf("Failed to get agent by name: %v", err)
	}
	if byName == nil {
		t.Fatal("Expected agent to exist by name")
	}

	// Get all agents
	all, err := store.GetAllAgents()
	if err != nil {
		t.Fatalf("Failed to get all agents: %v", err)
	}
	if len(all) != 1 {
		t.Errorf("Expected 1 agent, got %d", len(all))
	}
}

func TestAgent_UpdateLastSeen(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	_ = store.CreateOrUpdateAgent("SeenBot", "seen-uuid")

	// Update last seen
	if err := store.UpdateAgentLastSeen("seen-uuid"); err != nil {
		t.Fatalf("Failed to update last seen: %v", err)
	}

	agent, _ := store.GetAgentByAgentID("seen-uuid")
	if agent.LastSeen == nil {
		t.Error("LastSeen should be set after update")
	}
}

func TestAgent_TrustLevels(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	// Check trust for unknown agent (new)
	trust, err := store.CheckAgentTrust("UnknownBot", "unknown-uuid")
	if err != nil {
		t.Fatalf("Failed to check trust: %v", err)
	}
	if trust != models.AgentTrustNew {
		t.Errorf("Expected AgentTrustNew, got %v", trust)
	}

	// Create agent
	_ = store.CreateOrUpdateAgent("TrustBot", "trust-uuid")

	// Check trust for recognized agent
	trust, _ = store.CheckAgentTrust("TrustBot", "trust-uuid")
	if trust != models.AgentTrustRecognized {
		t.Errorf("Expected AgentTrustRecognized, got %v", trust)
	}

	// Check trust for suspicious agent (same name, different uuid)
	trust, _ = store.CheckAgentTrust("TrustBot", "different-uuid")
	if trust != models.AgentTrustSuspicious {
		t.Errorf("Expected AgentTrustSuspicious, got %v", trust)
	}
}

func TestAgent_Revoke(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	_ = store.CreateOrUpdateAgent("RevokeBot", "revoke-uuid")

	// Verify agent exists
	agent, _ := store.GetAgentByAgentID("revoke-uuid")
	if agent == nil {
		t.Fatal("Agent should exist")
	}

	// Revoke agent
	if err := store.RevokeAgent("revoke-uuid"); err != nil {
		t.Fatalf("Failed to revoke agent: %v", err)
	}

	// After revoke, agent should still exist but be in different state
	// (revoke doesn't delete, just marks somehow - let's verify it doesn't error)
}

func TestAgent_NonExistent(t *testing.T) {
	store := setupTestStoreWithAgents(t)
	defer func() { _ = store.Close() }()

	// Get non-existent agent
	agent, err := store.GetAgentByAgentID("does-not-exist")
	if err != nil {
		t.Fatalf("Should not error for non-existent agent: %v", err)
	}
	if agent != nil {
		t.Error("Agent should be nil for non-existent")
	}

	// Get non-existent by name
	byName, err := store.GetAgentByName("unknown-name")
	if err != nil {
		t.Fatalf("Should not error for non-existent name: %v", err)
	}
	if byName != nil {
		t.Error("Agent should be nil for non-existent name")
	}

	// Check trust for non-existent (should be new)
	trust, _ := store.CheckAgentTrust("UnknownBot", "unknown-uuid")
	if trust != models.AgentTrustNew {
		t.Errorf("Expected AgentTrustNew for unknown, got %v", trust)
	}
}