1 files changed, 371 insertions, 0 deletions
diff --git a/internal/store/agents_test.go b/internal/store/agents_test.go
new file mode 100644
index 0000000..1285e55
--- /dev/null
+++ b/internal/store/agents_test.go
@@ -0,0 +1,371 @@
+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)
+	}
+}