ITWorx-Pulse-Public/cmd/agent/agent_test.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/itworx/pulse/internal/agentstore"
	"github.com/itworx/pulse/internal/array"
	"github.com/itworx/pulse/internal/container"
	"github.com/itworx/pulse/internal/disk"
	"github.com/itworx/pulse/internal/host"
	"github.com/itworx/pulse/internal/pool"
	"github.com/itworx/pulse/internal/process"
	"github.com/itworx/pulse/internal/runtimeconfig"
	"github.com/itworx/pulse/internal/share"
)

// fakeWriter records what the agent publishes and can fail a chosen capability.
type fakeWriter struct {
	mu         sync.Mutex
	puts       []agentstore.Snapshot
	deadlines  []bool
	failures   map[agentstore.Capability]error
	blockUntil chan struct{}
}

func newFakeWriter() *fakeWriter {
	return &fakeWriter{failures: map[agentstore.Capability]error{}}
}

func (w *fakeWriter) Put(ctx context.Context, snapshot agentstore.Snapshot) error {
	w.mu.Lock()
	failure := w.failures[snapshot.Capability]
	block := w.blockUntil
	w.mu.Unlock()
	if block != nil {
		select {
		case <-block:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	_, hasDeadline := ctx.Deadline()
	w.mu.Lock()
	defer w.mu.Unlock()
	w.deadlines = append(w.deadlines, hasDeadline)
	if failure != nil {
		return failure
	}
	w.puts = append(w.puts, snapshot)
	return nil
}

func (w *fakeWriter) recorded() []agentstore.Snapshot {
	w.mu.Lock()
	defer w.mu.Unlock()
	return append([]agentstore.Snapshot(nil), w.puts...)
}

func (w *fakeWriter) countFor(capability agentstore.Capability) int {
	count := 0
	for _, snapshot := range w.recorded() {
		if snapshot.Capability == capability {
			count++
		}
	}
	return count
}

func (w *fakeWriter) failCapability(capability agentstore.Capability, err error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.failures[capability] = err
}

// fakeSource stands in for the procfs collector.
type fakeSource struct {
	mu          sync.Mutex
	hostErr     error
	processErr  error
	hostCalls   int
	procCalls   int
	release     chan struct{}
	hugePayload bool
}

type fakeContainerSource struct{ *fakeSource }

func (s fakeContainerSource) Containers(context.Context) (container.RawSnapshot, error) {
	return container.RawSnapshot{Source: container.Source{ID: "unraid", Type: "unraid"}, Containers: []container.RawContainer{{ID: "runtime-1", Name: "pulse-api", State: "running", Health: "healthy"}}, ObservedAt: time.Date(2026, 8, 10, 2, 0, 0, 0, time.UTC)}, nil
}

type fakeInventorySource struct{ *fakeSource }

func (s fakeInventorySource) Containers(context.Context) (container.RawSnapshot, error) {
	return container.RawSnapshot{Source: container.Source{ID: "unraid", Type: "unraid"}, Containers: []container.RawContainer{}, ObservedAt: time.Date(2026, 8, 10, 2, 0, 0, 0, time.UTC)}, nil
}
func (s fakeInventorySource) Array(context.Context) (array.RawSnapshot, error) {
	return array.RawSnapshot{Source: array.Source{ID: "unraid", Type: "unraid"}, State: array.StateOperational, Members: []array.RawMember{}, ObservedAt: time.Date(2026, 8, 10, 2, 0, 0, 0, time.UTC)}, nil
}
func (s fakeInventorySource) Disks(context.Context) (disk.RawSnapshot, error) {
	return disk.RawSnapshot{Source: disk.Source{ID: "unraid", Type: "unraid"}, Disks: []disk.RawDisk{}, ObservedAt: time.Date(2026, 8, 10, 2, 0, 0, 0, time.UTC)}, nil
}
func (s fakeInventorySource) Pools(context.Context) (pool.RawSnapshot, error) {
	return pool.RawSnapshot{Source: pool.Source{ID: "unraid", Type: "unraid"}, Pools: []pool.RawPool{}, ObservedAt: time.Date(2026, 8, 10, 2, 0, 0, 0, time.UTC)}, nil
}
func (s fakeInventorySource) Shares(context.Context) (share.RawSnapshot, error) {
	return share.RawSnapshot{Source: share.Source{ID: "unraid", Type: "unraid"}, Shares: []share.RawShare{}, ObservedAt: time.Date(2026, 8, 10, 2, 0, 0, 0, time.UTC)}, nil
}

func (s *fakeSource) Host(ctx context.Context) (host.RawSnapshot, error) {
	s.mu.Lock()
	release := s.release
	err := s.hostErr
	s.hostCalls++
	huge := s.hugePayload
	s.mu.Unlock()
	if release != nil {
		select {
		case <-release:
		case <-ctx.Done():
			return host.RawSnapshot{}, ctx.Err()
		}
	}
	if err != nil {
		return host.RawSnapshot{}, err
	}
	snapshot := host.RawSnapshot{
		Identity:      host.HostIdentity{Name: "tower"},
		UptimeSeconds: 100,
		Memory:        host.RawMemory{TotalBytes: 1024, AvailableBytes: 512},
		ObservedAt:    time.Date(2026, 8, 4, 12, 0, 0, 0, time.UTC),
	}
	if huge {
		warning := strings.Repeat("x", agentstore.MaxPayloadBytes)
		snapshot.Warnings = []string{warning}
	}
	return snapshot, nil
}

func (s *fakeSource) Processes(context.Context) (process.RawSnapshot, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.procCalls++
	if s.processErr != nil {
		return process.RawSnapshot{}, s.processErr
	}
	return process.RawSnapshot{
		Processes:  []process.RawProcess{{PID: 1, Name: "init", State: "sleeping"}},
		ObservedAt: time.Date(2026, 8, 4, 12, 0, 0, 0, time.UTC),
	}, nil
}

func (s *fakeSource) calls() (int, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.hostCalls, s.procCalls
}

type testClock struct {
	mu      sync.Mutex
	current time.Time
	// step advances the clock on every read. A loop test needs strictly increasing
	// timestamps, otherwise two iterations can observe the same instant and the second
	// one is legitimately "not due yet".
	step time.Duration
}

func (c *testClock) now() time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	current := c.current
	c.current = c.current.Add(c.step)
	return current
}

func (c *testClock) setStep(step time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.step = step
}

func (c *testClock) advance(d time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.current = c.current.Add(d)
}

func newTestAgent(t *testing.T, source snapshotSource, writer agentstore.Writer, collectInterval time.Duration) (*agent, *testClock, string) {
	t.Helper()
	heartbeat := filepath.Join(t.TempDir(), "healthy")
	config := runtimeconfig.AgentConfig{
		AgentID:         "pulse-agent-test",
		CollectInterval: collectInterval,
		Service:         runtimeconfig.ServiceConfig{ServiceName: "agent", HeartbeatFile: heartbeat},
	}
	logger := slog.New(slog.NewJSONHandler(io.Discard, nil))
	instance := newAgent(config, source, writer, logger)
	clock := &testClock{current: time.Date(2026, 8, 4, 12, 0, 0, 0, time.UTC)}
	instance.now = clock.now
	return instance, clock, heartbeat
}

func readHeartbeat(t *testing.T, path string) string {
	t.Helper()
	data, err := os.ReadFile(path) //nolint:gosec // test-controlled path
	if err != nil {
		t.Fatalf("read heartbeat: %v", err)
	}
	return strings.TrimSpace(string(data))
}

func TestIterateCollectsOnTheIntervalAndHeartbeatsEveryIteration(t *testing.T) {
	writer := newFakeWriter()
	source := &fakeSource{}
	// A 15 second collection interval must not stretch the loop: newAgent caps the
	// loop at maxLoopInterval so the heartbeat stays inside the contract's 10s window.
	instance, clock, heartbeat := newTestAgent(t, source, writer, 15*time.Second)
	if instance.loopInterval != maxLoopInterval {
		t.Fatalf("loop interval = %s, want %s", instance.loopInterval, maxLoopInterval)
	}

	ctx := context.Background()
	for iteration := 0; iteration < 7; iteration++ {
		instance.iterate(ctx)
		if got := readHeartbeat(t, heartbeat); got != clock.now().UTC().Format(time.RFC3339) {
			t.Fatalf("iteration %d heartbeat = %q, want the current time", iteration, got)
		}
		clock.advance(instance.loopInterval)
	}

	// Seven iterations five seconds apart cover t=0s..t=30s: collection is due at
	// t=0s, t=15s and t=30s. The four iterations in between do no work, yet every one
	// of them completed and refreshed the heartbeat, which is what the contract asks
	// for ("an empty poll is still a completed iteration").
	hostCalls, procCalls := source.calls()
	if hostCalls != 3 || procCalls != 3 {
		t.Fatalf("collections = host %d, processes %d; want 3 each over 30 seconds", hostCalls, procCalls)
	}
	if writer.countFor(agentstore.CapabilityHost) != 3 || writer.countFor(agentstore.CapabilityProcesses) != 3 {
		t.Fatalf("unexpected publications: %+v", writer.recorded())
	}
}

func TestOneFailingCapabilityDoesNotStopTheOthers(t *testing.T) {
	writer := newFakeWriter()
	source := &fakeSource{hostErr: errors.New("procfs read failed")}
	instance, _, _ := newTestAgent(t, source, writer, time.Second)

	instance.collectOnce(context.Background())

	recorded := writer.recorded()
	if len(recorded) != 1 || recorded[0].Capability != agentstore.CapabilityProcesses {
		t.Fatalf("expected only the process snapshot, got %+v", recorded)
	}

	// The reverse case: the store rejects one capability.
	writer.failCapability(agentstore.CapabilityProcesses, errors.New("store rejected"))
	source.mu.Lock()
	source.hostErr = nil
	source.mu.Unlock()
	instance.collectOnce(context.Background())

	if writer.countFor(agentstore.CapabilityHost) != 1 {
		t.Fatalf("host must still be published: %+v", writer.recorded())
	}
	if writer.countFor(agentstore.CapabilityProcesses) != 1 {
		t.Fatalf("the rejected capability must not be recorded twice: %+v", writer.recorded())
	}
}

func TestFailedCollectionWritesNothingRatherThanAnEmptySnapshot(t *testing.T) {
	writer := newFakeWriter()
	source := &fakeSource{hostErr: errors.New("boom"), processErr: errors.New("boom")}
	instance, _, _ := newTestAgent(t, source, writer, time.Second)

	instance.collectOnce(context.Background())

	if recorded := writer.recorded(); len(recorded) != 0 {
		t.Fatalf("a failed collection must publish nothing, got %+v", recorded)
	}
}

func TestOversizedPayloadIsRefused(t *testing.T) {
	writer := newFakeWriter()
	source := &fakeSource{hugePayload: true}
	instance, _, _ := newTestAgent(t, source, writer, time.Second)

	if _, err := instance.publish(context.Background(), instance.capabilities[0]); err == nil {
		t.Fatal("expected an oversized payload to be refused")
	}
	if writer.countFor(agentstore.CapabilityHost) != 0 {
		t.Fatal("an oversized payload must never reach the store")
	}
}

func TestPublishUsesABoundedContext(t *testing.T) {
	writer := newFakeWriter()
	instance, _, _ := newTestAgent(t, &fakeSource{}, writer, time.Second)

	instance.collectOnce(context.Background())

	writer.mu.Lock()
	defer writer.mu.Unlock()
	if len(writer.deadlines) == 0 {
		t.Fatal("no writes recorded")
	}
	for index, bounded := range writer.deadlines {
		if !bounded {
			t.Fatalf("write %d ran without a deadline", index)
		}
	}
}

func TestPublishedSnapshotCarriesTheObservedTimeAndPayload(t *testing.T) {
	writer := newFakeWriter()
	instance, _, _ := newTestAgent(t, &fakeSource{}, writer, time.Second)

	instance.collectOnce(context.Background())

	for _, snapshot := range writer.recorded() {
		if snapshot.AgentID != "pulse-agent-test" {
			t.Fatalf("agent id = %q", snapshot.AgentID)
		}
		if !snapshot.Capability.Valid() {
			t.Fatalf("unknown capability %q", snapshot.Capability)
		}
		if snapshot.ObservedAt.IsZero() || snapshot.ObservedAt.Location() != time.UTC {
			t.Fatalf("observed at = %v", snapshot.ObservedAt)
		}
		var decoded map[string]any
		if err := json.Unmarshal(snapshot.Payload, &decoded); err != nil {
			t.Fatalf("payload is not a JSON object: %v", err)
		}
	}
}

func TestRunWritesTheFirstHeartbeatBeforeTheFirstBlockingCall(t *testing.T) {
	writer := newFakeWriter()
	release := make(chan struct{})
	source := &fakeSource{release: release}
	instance, clock, heartbeat := newTestAgent(t, source, writer, time.Second)
	ticks := make(chan time.Time)
	instance.newTicker = func(time.Duration) (<-chan time.Time, func()) { return ticks, func() {} }

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- instance.run(ctx) }()

	// The collector is blocked inside the first iteration, yet the heartbeat must
	// already exist: a slow cold start is not a hang.
	deadline := time.Now().Add(2 * time.Second)
	for {
		if _, err := os.Stat(heartbeat); err == nil {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("no heartbeat was written before the first blocking call")
		}
		time.Sleep(time.Millisecond)
	}
	if got := readHeartbeat(t, heartbeat); got != clock.now().UTC().Format(time.RFC3339) {
		t.Fatalf("heartbeat = %q", got)
	}

	close(release)
	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("run returned error: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("run did not return after cancellation")
	}
}

func TestRunHeartbeatsFromTheLoopAndShutsDownCleanly(t *testing.T) {
	writer := newFakeWriter()
	source := &fakeSource{}
	instance, clock, heartbeat := newTestAgent(t, source, writer, time.Second)
	clock.setStep(time.Second)
	start := clock.now()
	ticks := make(chan time.Time)
	instance.newTicker = func(time.Duration) (<-chan time.Time, func()) { return ticks, func() {} }

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- instance.run(ctx) }()

	// An unbuffered tick is only received once the previous iteration has finished, so
	// each successful send proves one more completed loop iteration.
	for iteration := 0; iteration < 3; iteration++ {
		select {
		case ticks <- time.Now():
		case <-time.After(5 * time.Second):
			t.Fatalf("iteration %d never completed", iteration)
		}
	}
	if hostCalls, procCalls := source.calls(); hostCalls < 3 || procCalls < 3 {
		t.Fatalf("collections = host %d, processes %d; want at least 3 each", hostCalls, procCalls)
	}

	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("run returned error: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("run did not stop on cancellation")
	}

	// Read the heartbeat only after the loop has stopped: the file is truncated and
	// rewritten in place, so a concurrent reader can legitimately observe it empty.
	// The healthcheck script reads only the mtime, which is why that is safe.
	written, err := time.Parse(time.RFC3339, readHeartbeat(t, heartbeat))
	if err != nil {
		t.Fatalf("heartbeat content is not an RFC 3339 timestamp: %v", err)
	}
	if written.Before(start) {
		t.Fatalf("heartbeat %s predates the loop start %s", written, start)
	}
}

func TestRunDoesNotBlockShutdownOnAStuckWrite(t *testing.T) {
	writer := newFakeWriter()
	writer.blockUntil = make(chan struct{})
	instance, _, _ := newTestAgent(t, &fakeSource{}, writer, time.Second)
	ticks := make(chan time.Time)
	instance.newTicker = func(time.Duration) (<-chan time.Time, func()) { return ticks, func() {} }

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- instance.run(ctx) }()

	time.Sleep(20 * time.Millisecond)
	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("run returned error: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("a blocked store write must not hold up shutdown")
	}
}

func TestRunRefusesToStartWhenTheHelloIsNotValid(t *testing.T) {
	writer := newFakeWriter()
	instance, _, _ := newTestAgent(t, &fakeSource{}, writer, time.Second)
	instance.agentID = ""

	if err := instance.run(context.Background()); err == nil {
		t.Fatal("an agent whose own hello fails validation must refuse to start")
	}
	if len(writer.recorded()) != 0 {
		t.Fatal("nothing may be published before the self-check passes")
	}
}

func TestRunRequiresAWriter(t *testing.T) {
	instance, _, _ := newTestAgent(t, &fakeSource{}, nil, time.Second)
	instance.writer = nil
	if err := instance.run(context.Background()); !errors.Is(err, errNoWriter) {
		t.Fatalf("error = %v, want errNoWriter", err)
	}
}

func TestHelloAnnouncesOnlyReadOnlyCapabilities(t *testing.T) {
	instance, clock, _ := newTestAgent(t, &fakeSource{}, newFakeWriter(), time.Second)
	hello := instance.hello()
	if err := hello.Validate(clock.now()); err != nil {
		t.Fatalf("hello validation failed: %v", err)
	}
	if len(hello.Capabilities) != 2 {
		t.Fatalf("capabilities = %+v", hello.Capabilities)
	}
	for _, item := range hello.Capabilities {
		if !item.ReadOnly {
			t.Fatalf("capability %q is not read-only", item.ID)
		}
		if !agentstore.Capability(item.ID).Valid() {
			t.Fatalf("capability %q is not a recognised store capability", item.ID)
		}
	}
}

func TestContainerCapabilityIsPublishedOnlyWhenSourceSupportsIt(t *testing.T) {
	writer := newFakeWriter()
	instance, _, _ := newTestAgent(t, fakeContainerSource{fakeSource: &fakeSource{}}, writer, time.Second)
	if len(instance.hello().Capabilities) != 3 {
		t.Fatalf("capabilities = %+v, want host/processes/containers", instance.hello().Capabilities)
	}
	instance.collectOnce(context.Background())
	if writer.countFor(agentstore.CapabilityContainers) != 1 {
		t.Fatalf("container snapshot was not published: %+v", writer.recorded())
	}
}

func TestInventoryCapabilitiesArePublishedOnlyForConfiguredUnraidSource(t *testing.T) {
	writer := newFakeWriter()
	instance, _, _ := newTestAgent(t, fakeInventorySource{fakeSource: &fakeSource{}}, writer, time.Second)
	if got := len(instance.hello().Capabilities); got != 7 {
		t.Fatalf("capabilities = %+v, want host/processes/containers/array/disks/pools/shares", instance.hello().Capabilities)
	}
	instance.collectOnce(context.Background())
	for _, capability := range []agentstore.Capability{agentstore.CapabilityContainers, agentstore.CapabilityArray, agentstore.CapabilityDisks, agentstore.CapabilityPools, agentstore.CapabilityShares} {
		if writer.countFor(capability) != 1 {
			t.Fatalf("capability %q was not published: %+v", capability, writer.recorded())
		}
	}
}

func TestHeartbeatFailureIsNotFatal(t *testing.T) {
	writer := newFakeWriter()
	instance, _, _ := newTestAgent(t, &fakeSource{}, writer, time.Second)
	instance.heartbeatPath = filepath.Join(t.TempDir(), "missing-directory", "healthy")

	instance.iterate(context.Background())

	if len(writer.recorded()) == 0 {
		t.Fatal("a heartbeat write failure must not stop the collection loop")
	}
}

func TestSnapshotWriterRefusesAMissingPool(t *testing.T) {
	// A nil pool would let every publish fail deep in the driver. The agent must refuse
	// to start instead, so an operator sees the misconfiguration rather than an agent
	// that appears healthy while storing nothing.
	if _, err := newSnapshotWriter(nil); !errors.Is(err, errStoreNotLinked) {
		t.Fatalf("error = %v, want errStoreNotLinked", err)
	}
}