CHORUS/pkg/election/election_test.go

package election

import (
	"context"
	"testing"
	"time"

	"chorus/pkg/config"
)

func TestElectionManager_NewElectionManager(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	if em == nil {
		t.Fatal("Expected NewElectionManager to return non-nil manager")
	}

	if em.nodeID != "test-node" {
		t.Errorf("Expected nodeID to be 'test-node', got %s", em.nodeID)
	}

	if em.state != StateIdle {
		t.Errorf("Expected initial state to be StateIdle, got %v", em.state)
	}
}

func TestElectionManager_StartElection(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Start election
	err := em.StartElection()
	if err != nil {
		t.Fatalf("Failed to start election: %v", err)
	}

	// Verify state changed
	if em.state != StateCandidate {
		t.Errorf("Expected state to be StateCandidate after starting election, got %v", em.state)
	}

	// Verify we added ourselves as a candidate
	em.mu.RLock()
	candidate, exists := em.candidates[em.nodeID]
	em.mu.RUnlock()

	if !exists {
		t.Error("Expected to find ourselves as a candidate after starting election")
	}

	if candidate.NodeID != em.nodeID {
		t.Errorf("Expected candidate NodeID to be %s, got %s", em.nodeID, candidate.NodeID)
	}
}

func TestElectionManager_Vote(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Add a candidate first
	candidate := &AdminCandidate{
		NodeID:      "candidate-1",
		Term:        1,
		Score:       0.8,
		Capabilities: []string{"admin"},
		LastSeen:    time.Now(),
	}
	
	em.mu.Lock()
	em.candidates["candidate-1"] = candidate
	em.mu.Unlock()

	// Vote for the candidate
	err := em.Vote("candidate-1")
	if err != nil {
		t.Fatalf("Failed to vote: %v", err)
	}

	// Verify vote was recorded
	em.mu.RLock()
	vote, exists := em.votes[em.nodeID]
	em.mu.RUnlock()

	if !exists {
		t.Error("Expected to find our vote after voting")
	}

	if vote != "candidate-1" {
		t.Errorf("Expected vote to be for 'candidate-1', got %s", vote)
	}
}

func TestElectionManager_VoteInvalidCandidate(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Try to vote for non-existent candidate
	err := em.Vote("non-existent")
	if err == nil {
		t.Error("Expected error when voting for non-existent candidate")
	}
}

func TestElectionManager_AddCandidate(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	candidate := &AdminCandidate{
		NodeID:      "new-candidate",
		Term:        1,
		Score:       0.7,
		Capabilities: []string{"admin", "leader"},
		LastSeen:    time.Now(),
	}

	err := em.AddCandidate(candidate)
	if err != nil {
		t.Fatalf("Failed to add candidate: %v", err)
	}

	// Verify candidate was added
	em.mu.RLock()
	stored, exists := em.candidates["new-candidate"]
	em.mu.RUnlock()

	if !exists {
		t.Error("Expected to find added candidate")
	}

	if stored.NodeID != "new-candidate" {
		t.Errorf("Expected stored candidate NodeID to be 'new-candidate', got %s", stored.NodeID)
	}

	if stored.Score != 0.7 {
		t.Errorf("Expected stored candidate score to be 0.7, got %f", stored.Score)
	}
}

func TestElectionManager_FindElectionWinner(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Add candidates with different scores
	candidates := []*AdminCandidate{
		{
			NodeID:      "candidate-1",
			Term:        1,
			Score:       0.6,
			Capabilities: []string{"admin"},
			LastSeen:    time.Now(),
		},
		{
			NodeID:      "candidate-2", 
			Term:        1,
			Score:       0.8,
			Capabilities: []string{"admin", "leader"},
			LastSeen:    time.Now(),
		},
		{
			NodeID:      "candidate-3",
			Term:        1,
			Score:       0.7,
			Capabilities: []string{"admin"},
			LastSeen:    time.Now(),
		},
	}

	em.mu.Lock()
	for _, candidate := range candidates {
		em.candidates[candidate.NodeID] = candidate
	}
	
	// Add some votes
	em.votes["voter-1"] = "candidate-2"
	em.votes["voter-2"] = "candidate-2" 
	em.votes["voter-3"] = "candidate-1"
	em.mu.Unlock()

	// Find winner
	winner := em.findElectionWinner()
	
	if winner == nil {
		t.Fatal("Expected findElectionWinner to return a winner")
	}

	// candidate-2 should win with most votes (2 votes)
	if winner.NodeID != "candidate-2" {
		t.Errorf("Expected winner to be 'candidate-2', got %s", winner.NodeID)
	}
}

func TestElectionManager_FindElectionWinnerNoVotes(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Add candidates but no votes - should fall back to highest score
	candidates := []*AdminCandidate{
		{
			NodeID:      "candidate-1",
			Term:        1,
			Score:       0.6,
			Capabilities: []string{"admin"},
			LastSeen:    time.Now(),
		},
		{
			NodeID:      "candidate-2",
			Term:        1,
			Score:       0.9, // Highest score
			Capabilities: []string{"admin", "leader"},
			LastSeen:    time.Now(),
		},
	}

	em.mu.Lock()
	for _, candidate := range candidates {
		em.candidates[candidate.NodeID] = candidate
	}
	em.mu.Unlock()

	// Find winner without any votes
	winner := em.findElectionWinner()
	
	if winner == nil {
		t.Fatal("Expected findElectionWinner to return a winner")
	}

	// candidate-2 should win with highest score
	if winner.NodeID != "candidate-2" {
		t.Errorf("Expected winner to be 'candidate-2' (highest score), got %s", winner.NodeID)
	}
}

func TestElectionManager_HandleElectionVote(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Add a candidate first
	candidate := &AdminCandidate{
		NodeID:      "candidate-1",
		Term:        1,
		Score:       0.8,
		Capabilities: []string{"admin"},
		LastSeen:    time.Now(),
	}
	
	em.mu.Lock()
	em.candidates["candidate-1"] = candidate
	em.mu.Unlock()

	// Create vote message
	msg := ElectionMessage{
		Type:   MessageTypeVote,
		NodeID: "voter-1",
		Data: map[string]interface{}{
			"candidate": "candidate-1",
		},
	}

	// Handle the vote
	em.handleElectionVote(msg)

	// Verify vote was recorded
	em.mu.RLock()
	vote, exists := em.votes["voter-1"]
	em.mu.RUnlock()

	if !exists {
		t.Error("Expected vote to be recorded after handling vote message")
	}

	if vote != "candidate-1" {
		t.Errorf("Expected recorded vote to be for 'candidate-1', got %s", vote)
	}
}

func TestElectionManager_HandleElectionVoteInvalidData(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Create vote message with invalid data
	msg := ElectionMessage{
		Type:   MessageTypeVote,
		NodeID: "voter-1",
		Data:   "invalid-data", // Should be map[string]interface{}
	}

	// Handle the vote - should not crash
	em.handleElectionVote(msg)

	// Verify no vote was recorded
	em.mu.RLock()
	_, exists := em.votes["voter-1"]
	em.mu.RUnlock()

	if exists {
		t.Error("Expected no vote to be recorded with invalid data")
	}
}

func TestElectionManager_CompleteElection(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Set up election state
	em.mu.Lock()
	em.state = StateCandidate
	em.currentTerm = 1
	em.mu.Unlock()

	// Add a candidate
	candidate := &AdminCandidate{
		NodeID:      "winner",
		Term:        1,
		Score:       0.9,
		Capabilities: []string{"admin", "leader"},
		LastSeen:    time.Now(),
	}
	
	em.mu.Lock()
	em.candidates["winner"] = candidate
	em.mu.Unlock()

	// Complete election
	em.CompleteElection()

	// Verify state reset
	em.mu.RLock()
	state := em.state
	em.mu.RUnlock()

	if state != StateIdle {
		t.Errorf("Expected state to be StateIdle after completing election, got %v", state)
	}
}

func TestElectionManager_Concurrency(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-node",
		},
	}

	em := NewElectionManager(cfg)
	
	// Test concurrent access to vote and candidate operations
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	// Add a candidate
	candidate := &AdminCandidate{
		NodeID:      "candidate-1",
		Term:        1,
		Score:       0.8,
		Capabilities: []string{"admin"},
		LastSeen:    time.Now(),
	}
	
	err := em.AddCandidate(candidate)
	if err != nil {
		t.Fatalf("Failed to add candidate: %v", err)
	}

	// Run concurrent operations
	done := make(chan bool, 2)

	// Concurrent voting
	go func() {
		defer func() { done <- true }()
		for i := 0; i < 10; i++ {
			select {
			case <-ctx.Done():
				return
			default:
				em.Vote("candidate-1") // Ignore errors in concurrent test
				time.Sleep(10 * time.Millisecond)
			}
		}
	}()

	// Concurrent state checking
	go func() {
		defer func() { done <- true }()
		for i := 0; i < 10; i++ {
			select {
			case <-ctx.Done():
				return
			default:
				em.findElectionWinner() // Just check for races
				time.Sleep(10 * time.Millisecond)
			}
		}
	}()

	// Wait for completion
	for i := 0; i < 2; i++ {
		select {
		case <-done:
		case <-ctx.Done():
			t.Fatal("Concurrent test timed out")
		}
	}
}