orca/tests/integration/harness.go

// Package integration provides the hermetic test harness for orca
// integration tests (P08, REQ-087). The harness creates a temp
// ORCA_HOME, initializes a cluster (DB + CA + master key), registers
// mock peers that record SSH commands without a real SSH server, and
// exposes helpers (SubmitJob, ListAllocs, DrainNode, VerifyState) for
// the scenario tests.
//
// The harness is hermetic: every test gets its own temp ORCA_HOME
// (t.Setenv), so there is no cross-test state leakage and no reliance
// on the operator's ~/.orca. Mock peers record every Exec call so the
// drain/migrate scenarios can assert on the exact systemctl commands
// issued without a real systemd host.
package integration

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/google/uuid"

	"git.cloudinit.dev/coreci/orca/internal/certpaths"
	"git.cloudinit.dev/coreci/orca/internal/engine"
	"git.cloudinit.dev/coreci/orca/internal/model"
	"git.cloudinit.dev/coreci/orca/internal/paths"
	"git.cloudinit.dev/coreci/orca/internal/secrets"
	"git.cloudinit.dev/coreci/orca/internal/security"
	"git.cloudinit.dev/coreci/orca/internal/store"
)

// Peer is a mock SSH peer registered with the harness. It records
// every Exec command the harness issues against it so scenarios can
// assert on the exact systemctl invocations without a real SSH server.
// RunningAllocs is the set of orca-alloc-<id>.service units the mock
// reports as running; listRunningAllocs-style commands read it, and
// stop/start commands mutate it.
type Peer struct {
	Name     string
	Address  string
	NodeID   string
	mu       sync.Mutex
	commands []string
	running  map[string]bool
}

// Commands returns a snapshot of the recorded Exec commands.
func (p *Peer) Commands() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := make([]string, len(p.commands))
	copy(out, p.commands)
	return out
}

// Exec records cmd and simulates a systemd host:
//   - "systemctl list-units 'orca-alloc-*.service' ..." reports the
//     currently-running alloc ids.
//   - "systemctl stop orca-alloc-<id>.service" removes <id> from the
//     running set.
//   - "systemctl start orca-alloc-<id>.service" adds <id> to the
//     running set.
//   - everything else is recorded and returns empty output.
func (p *Peer) Exec(_ context.Context, _ string, cmd string) ([]byte, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.commands = append(p.commands, cmd)
	switch {
	case startsWith(cmd, "systemctl list-units 'orca-alloc-*.service'"):
		var lines []byte
		for id := range p.running {
			lines = append(lines, []byte(fmt.Sprintf("orca-alloc-%s.service loaded active running orca alloc %s\n", id, id))...)
		}
		return lines, nil
	case startsWith(cmd, "systemctl stop orca-alloc-"):
		id := allocIDFromStopCmd(cmd)
		delete(p.running, id)
		return nil, nil
	case startsWith(cmd, "systemctl start orca-alloc-"):
		id := allocIDFromStartCmd(cmd)
		if id != "" {
			p.running[id] = true
		}
		return nil, nil
	default:
		return nil, nil
	}
}

// RunningAllocs returns the currently-running alloc ids (snapshot).
func (p *Peer) RunningAllocs() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := make([]string, 0, len(p.running))
	for id := range p.running {
		out = append(out, id)
	}
	return out
}

// SetRunning injects alloc ids into the running set (for scenarios that
// need a peer to start with running allocations before a drain).
func (p *Peer) SetRunning(ids ...string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, id := range ids {
		if p.running == nil {
			p.running = make(map[string]bool)
		}
		p.running[id] = true
	}
}

// MockTransport is the harness's drainExecer: it routes Exec calls to
// the registered Peer with the matching name/address. It satisfies
// the cli.drainExecer interface (Exec(ctx, peer, cmd) ([]byte, error))
// without importing the cli package (which would create an import
// cycle).
type MockTransport struct {
	mu    sync.RWMutex
	peers map[string]*Peer
}

// NewMockTransport returns an empty MockTransport.
func NewMockTransport() *MockTransport {
	return &MockTransport{peers: make(map[string]*Peer)}
}

// Register associates peer with the harness transport under both the
// peer's Name and Address so peerAddrForNode-style "host:22" lookups
// resolve.
func (t *MockTransport) Register(p *Peer) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.peers[p.Name] = p
	t.peers[p.Address] = p
}

// Exec routes cmd to the registered peer. Unknown peer returns an
// error so scenarios catch a missing registration.
func (t *MockTransport) Exec(ctx context.Context, peer string, cmd string) ([]byte, error) {
	t.mu.RLock()
	p, ok := t.peers[peer]
	t.mu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("mock transport: unknown peer %q", peer)
	}
	return p.Exec(ctx, peer, cmd)
}

// Harness is the hermetic test harness. Each test constructs one with
// NewHarness and defers h.Close to clean up the temp ORCA_HOME.
type Harness struct {
	t         *testing.T
	Home      string
	DB        *sql.DB
	Nodes     *store.NodeRepo
	Jobs      *store.JobRepo
	Tasks     *store.TaskRepo
	Registry  *engine.NodeRegistry
	Transport *MockTransport

	dbClose func() error
}

// NewHarness creates a temp ORCA_HOME, initializes the cluster (DB
// migrations + CA + master key + localhost node), and returns a ready
// harness. The temp dir is cleaned up via t.Cleanup.
func NewHarness(t *testing.T) *Harness {
	t.Helper()
	home := t.TempDir()
	t.Setenv("ORCA_HOME", home)

	h := &Harness{
		t:         t,
		Home:      home,
		Transport: NewMockTransport(),
	}

	if err := os.MkdirAll(paths.ClusterDir(), 0o755); err != nil {
		t.Fatalf("mkdir cluster dir: %v", err)
	}

	db, err := store.Open(certpaths.DBPath())
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	h.DB = db
	h.dbClose = db.Close
	h.Nodes = store.NewNodeRepo(db)
	h.Jobs = store.NewJobRepo(db)
	h.Tasks = store.NewTaskRepo(db)
	audit := engine.NewAudit(store.NewAuditRepo(db), testLogger())
	h.Registry = engine.NewNodeRegistry(h.Nodes, audit, testLogger())

	if _, err := security.CAInit(home, "orca-test-ca"); err != nil {
		t.Fatalf("CAInit: %v", err)
	}

	mk, err := secrets.GenerateMasterKey()
	if err != nil {
		t.Fatalf("GenerateMasterKey: %v", err)
	}
	if err := secrets.SaveMasterKey(paths.MasterKeyPath(), mk); err != nil {
		t.Fatalf("SaveMasterKey: %v", err)
	}

	ctx := context.Background()
	localhost := &model.Node{
		ID:       uuid.NewString(),
		Name:     "localhost",
		Address:  "localhost:8443",
		State:    model.NodeStateReady,
		JoinedAt: time.Now().UTC(),
		LastSeen: time.Now().UTC(),
		Kind:     string(model.NodeKindLocalhost),
		OS:       "linux",
	}
	if err := h.Registry.Join(ctx, localhost); err != nil {
		t.Fatalf("register localhost: %v", err)
	}

	t.Cleanup(func() { _ = h.Close() })
	return h
}

// RegisterPeer creates a mock Peer, registers it as a node in the
// cluster registry, and wires it into the mock transport so drain/
// migrate commands route to it. The peer starts with no running
// allocations; use Peer.SetRunning to seed it.
func (h *Harness) RegisterPeer(name, address string) *Peer {
	h.t.Helper()
	p := &Peer{
		Name:    name,
		Address: address,
		NodeID:  uuid.NewString(),
		running: make(map[string]bool),
	}
	ctx := context.Background()
	node := &model.Node{
		ID:       p.NodeID,
		Name:     name,
		Address:  address,
		State:    model.NodeStateReady,
		JoinedAt: time.Now().UTC(),
		LastSeen: time.Now().UTC(),
		Kind:     string(model.NodeKindLinux),
		OS:       "linux",
	}
	if err := h.Registry.Join(ctx, node); err != nil {
		h.t.Fatalf("register peer %q: %v", name, err)
	}
	h.Transport.Register(p)
	return p
}

// SubmitJob runs a single-task job on the local executor and returns
// the job ID. The command must be a real binary path resolvable on the
// test host (e.g. "/bin/true", "/bin/echo").
func (h *Harness) SubmitJob(ctx context.Context, name, command string, args ...string) (string, error) {
	h.t.Helper()
	exec := engine.NewExecutor(h.Jobs, h.Tasks, testLogger())
	job := &model.Job{
		ID:     uuid.NewString(),
		Name:   name,
		Spec:   command,
		Status: model.JobStatusPending,
	}
	specs := []engine.TaskSpec{{
		Name:    name,
		Command: command,
		Args:    args,
	}}
	if err := exec.Run(ctx, job, specs); err != nil {
		return job.ID, err
	}
	return job.ID, nil
}

// ListAllocs returns the currently-running alloc ids on the named
// peer (as reported by the mock transport). For the local node the
// harness has no mock systemd host, so it returns an empty slice.
func (h *Harness) ListAllocs(ctx context.Context, peerName string) ([]string, error) {
	h.t.Helper()
	p, ok := h.lookupPeer(peerName)
	if !ok {
		return nil, fmt.Errorf("ListAllocs: peer %q not registered", peerName)
	}
	cmd := "systemctl list-units 'orca-alloc-*.service' --type=service --state=running --no-legend --no-pager"
	if _, err := h.Transport.Exec(ctx, peerName, cmd); err != nil {
		return nil, err
	}
	return p.RunningAllocs(), nil
}

// DrainNode marks the named peer's node "draining", stops every
// running allocation on it via the mock transport, and marks it
// "drained". It mirrors the `orca node drain` flow without going
// through the CLI cobra rootCmd (so the harness stays hermetic and
// does not leak global flag state between scenarios).
func (h *Harness) DrainNode(ctx context.Context, peerName string) ([]string, error) {
	h.t.Helper()
	p, ok := h.lookupPeer(peerName)
	if !ok {
		return nil, fmt.Errorf("DrainNode: peer %q not registered", peerName)
	}
	nodes, err := h.Registry.List(ctx)
	if err != nil {
		return nil, err
	}
	var node *model.Node
	for _, n := range nodes {
		if n.Name == peerName {
			node = n
			break
		}
	}
	if node == nil {
		return nil, fmt.Errorf("DrainNode: node %q not in registry", peerName)
	}
	if err := h.Registry.SetNodeState(ctx, node.ID, string(model.NodeStateDraining)); err != nil {
		return nil, err
	}
	ids := p.RunningAllocs()
	var stopped []string
	for _, id := range ids {
		stopCmd := fmt.Sprintf("systemctl stop orca-alloc-%s.service", id)
		if _, err := h.Transport.Exec(ctx, peerName, stopCmd); err != nil {
			return stopped, err
		}
		stopped = append(stopped, id)
	}
	if err := h.Registry.SetNodeState(ctx, node.ID, string(model.NodeStateDrained)); err != nil {
		return stopped, err
	}
	return stopped, nil
}

// VerifyState asserts the harness's on-disk state is self-consistent:
// the cluster dir exists, the master key is present with mode 0600,
// and the SQLite DB opens. Scenarios call it after a backup/restore
// round-trip to confirm the restored tree is usable.
func (h *Harness) VerifyState() error {
	if info, err := os.Stat(paths.ClusterDir()); err != nil || !info.IsDir() {
		return fmt.Errorf("cluster dir missing: %v", err)
	}
	mkInfo, err := os.Stat(paths.MasterKeyPath())
	if err != nil {
		return fmt.Errorf("master key missing: %w", err)
	}
	if mkInfo.Mode().Perm() != secrets.MasterKeyMode {
		return fmt.Errorf("master key mode %04o, want %04o", mkInfo.Mode().Perm(), secrets.MasterKeyMode)
	}
	db, err := store.Open(certpaths.DBPath())
	if err != nil {
		return fmt.Errorf("db open: %w", err)
	}
	defer db.Close()
	if err := db.Ping(); err != nil {
		return fmt.Errorf("db ping: %w", err)
	}
	return nil
}

// MasterKeyPath returns the cluster master key path (paths.MasterKeyPath()).
func (h *Harness) MasterKeyPath() string { return paths.MasterKeyPath() }

// Root returns the temp ORCA_HOME root.
func (h *Harness) Root() string { return paths.Root() }

// Close releases the DB handle. It is idempotent and safe to call
// from t.Cleanup.
func (h *Harness) Close() error {
	if h.dbClose != nil {
		err := h.dbClose()
		h.dbClose = nil
		return err
	}
	return nil
}

// lookupPeer returns the registered Peer by name.
func (h *Harness) lookupPeer(name string) (*Peer, bool) {
	h.Transport.mu.RLock()
	defer h.Transport.mu.RUnlock()
	p, ok := h.Transport.peers[name]
	return p, ok
}

// testLogger returns a slog logger writing JSON to stderr at Info
// level. It mirrors internal/cli.newLogger so the harness logs the
// same way the CLI does.
func testLogger() *slog.Logger {
	return slog.New(slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo}))
}

// startsWith reports whether s begins with prefix.
func startsWith(s, prefix string) bool {
	return len(s) >= len(prefix) && s[:len(prefix)] == prefix
}

// allocIDFromStopCmd extracts the alloc id from a
// "systemctl stop orca-alloc-<id>.service" command.
func allocIDFromStopCmd(cmd string) string {
	const prefix = "systemctl stop orca-alloc-"
	const suffix = ".service"
	if !startsWith(cmd, prefix) {
		return ""
	}
	s := cmd[len(prefix):]
	s = trimSuffix(s, suffix)
	return s
}

// allocIDFromStartCmd extracts the alloc id from a
// "systemctl start orca-alloc-<id>.service" command.
func allocIDFromStartCmd(cmd string) string {
	const prefix = "systemctl start orca-alloc-"
	const suffix = ".service"
	if !startsWith(cmd, prefix) {
		return ""
	}
	s := cmd[len(prefix):]
	s = trimSuffix(s, suffix)
	return s
}

// trimSuffix removes a trailing suffix from s if present.
func trimSuffix(s, suffix string) string {
	if len(s) >= len(suffix) && s[len(s)-len(suffix):] == suffix {
		return s[:len(s)-len(suffix)]
	}
	return s
}

// Compile-time assertion that *MockTransport satisfies the drainExecer
// shape used by internal/cli (Exec(ctx, peer, cmd) ([]byte, error)).
// We can't import internal/cli (import cycle), so the assertion is
// structural via an anonymous interface.
var _ drainExecerShape = (*MockTransport)(nil)

type drainExecerShape interface {
	Exec(context.Context, string, string) ([]byte, error)
}

// statPerm returns the file permissions for path, surfacing the
// stat error so tests can assert on file mode without rewriting the
// os.Stat dance.
func statPerm(path string) (os.FileMode, error) {
	info, err := os.Stat(path)
	if err != nil {
		return 0, err
	}
	return info.Mode().Perm(), nil
}