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