CHORUS/pkg/election/election_test.go

package election

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"chorus/pkg/config"
	pubsubpkg "chorus/pubsub"
	libp2p "github.com/libp2p/go-libp2p"
)

// newTestElectionManager wires a real libp2p host and PubSub instance so the
// election manager exercises the same code paths used in production.
func newTestElectionManager(t *testing.T) *ElectionManager {
	t.Helper()

	ctx, cancel := context.WithCancel(context.Background())

	host, err := libp2p.New(libp2p.ListenAddrStrings("/ip4/127.0.0.1/tcp/0"))
	if err != nil {
		cancel()
		t.Fatalf("failed to create libp2p host: %v", err)
	}

	ps, err := pubsubpkg.NewPubSub(ctx, host, "", "")
	if err != nil {
		host.Close()
		cancel()
		t.Fatalf("failed to create pubsub: %v", err)
	}

	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID:             host.ID().String(),
			Role:           "context_admin",
			Capabilities:   []string{"admin_election", "context_curation"},
			Models:         []string{"meta/llama-3.1-8b-instruct"},
			Specialization: "coordination",
		},
		Security: config.SecurityConfig{},
	}

	em := NewElectionManager(ctx, cfg, host, ps, host.ID().String())

	t.Cleanup(func() {
		em.Stop()
		ps.Close()
		host.Close()
		cancel()
	})

	return em
}

func TestNewElectionManagerInitialState(t *testing.T) {
	em := newTestElectionManager(t)

	if em.state != StateIdle {
		t.Fatalf("expected initial state %q, got %q", StateIdle, em.state)
	}

	if em.currentTerm != 0 {
		t.Fatalf("expected initial term 0, got %d", em.currentTerm)
	}

	if em.nodeID == "" {
		t.Fatal("expected nodeID to be populated")
	}
}

func TestElectionManagerCanBeAdmin(t *testing.T) {
	em := newTestElectionManager(t)

	if !em.canBeAdmin() {
		t.Fatal("expected node to qualify for admin election")
	}

	em.config.Agent.Capabilities = []string{"runtime_support"}
	if em.canBeAdmin() {
		t.Fatal("expected node without admin capabilities to be ineligible")
	}
}

func TestFindElectionWinnerPrefersVotesThenScore(t *testing.T) {
	em := newTestElectionManager(t)

	em.mu.Lock()
	em.candidates = map[string]*AdminCandidate{
		"candidate-1": {
			NodeID: "candidate-1",
			PeerID: em.host.ID(),
			Score:  0.65,
		},
		"candidate-2": {
			NodeID: "candidate-2",
			PeerID: em.host.ID(),
			Score:  0.80,
		},
	}
	em.votes = map[string]string{
		"voter-a": "candidate-1",
		"voter-b": "candidate-2",
		"voter-c": "candidate-2",
	}
	em.mu.Unlock()

	winner := em.findElectionWinner()
	if winner == nil {
		t.Fatal("expected a winner to be selected")
	}
	if winner.NodeID != "candidate-2" {
		t.Fatalf("expected candidate-2 to win, got %s", winner.NodeID)
	}
}

func TestHandleElectionMessageAddsCandidate(t *testing.T) {
	em := newTestElectionManager(t)

	em.mu.Lock()
	em.currentTerm = 3
	em.state = StateElecting
	em.mu.Unlock()

	candidate := &AdminCandidate{
		NodeID:       "peer-2",
		PeerID:       em.host.ID(),
		Capabilities: []string{"admin_election"},
		Uptime:       time.Second,
		Score:        0.75,
	}

	payload, err := json.Marshal(candidate)
	if err != nil {
		t.Fatalf("failed to marshal candidate: %v", err)
	}

	var data map[string]interface{}
	if err := json.Unmarshal(payload, &data); err != nil {
		t.Fatalf("failed to unmarshal candidate payload: %v", err)
	}

	msg := ElectionMessage{
		Type:      "candidacy_announcement",
		NodeID:    "peer-2",
		Timestamp: time.Now(),
		Term:      3,
		Data:      data,
	}

	serialized, err := json.Marshal(msg)
	if err != nil {
		t.Fatalf("failed to marshal election message: %v", err)
	}

	em.handleElectionMessage(serialized)

	em.mu.RLock()
	_, exists := em.candidates["peer-2"]
	em.mu.RUnlock()

	if !exists {
		t.Fatal("expected candidacy announcement to register candidate")
	}
}

func TestSendAdminHeartbeatRequiresLeadership(t *testing.T) {
	em := newTestElectionManager(t)

	if err := em.SendAdminHeartbeat(); err == nil {
		t.Fatal("expected error when non-admin sends heartbeat")
	}

	if err := em.Start(); err != nil {
		t.Fatalf("failed to start election manager: %v", err)
	}

	em.mu.Lock()
	em.currentAdmin = em.nodeID
	em.mu.Unlock()

	if err := em.SendAdminHeartbeat(); err != nil {
		t.Fatalf("expected heartbeat to succeed for current admin, got error: %v", err)
	}
}