monlet/agent/internal/app/app_test.go

package app

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/monlet/agent/internal/client"
	"github.com/monlet/agent/internal/config"
	"github.com/monlet/agent/internal/metrics"
	"github.com/monlet/agent/internal/spool"
)

func TestPrometheusOnlyDoesNotPush(t *testing.T) {
	var hits int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&hits, 1)
		w.WriteHeader(http.StatusAccepted)
	}))
	defer srv.Close()

	dir := t.TempDir()
	cfg := &config.Config{
		AgentID:  "a1",
		Mode:     "prometheus_only",
		StateDir: dir,
		Server:   config.ServerConfig{URL: srv.URL, Token: "t", HeartbeatInterval: config.Duration{Duration: 50 * time.Millisecond}, BatchInterval: config.Duration{Duration: 50 * time.Millisecond}},
		Metrics:  config.MetricsConfig{Enabled: false},
		Checks: []config.CheckConfig{{
			ID:                "c1",
			Command:           []string{"sh", "-c", "echo hi"},
			Interval:          config.Duration{Duration: 100 * time.Millisecond},
			Timeout:           config.Duration{Duration: 50 * time.Millisecond},
			NotificationOwner: "server",
		}},
	}
	sp, _ := spool.Open(filepath.Join(dir, "spool"), 100, 1<<20)
	a := &App{
		Cfg: cfg, AgentID: "a1", Version: "t",
		Client: client.New(srv.URL, "t", "a1", "t"),
		Spool:  sp, Metrics: metrics.New(),
		Log: slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 400*time.Millisecond)
	defer cancel()
	_ = a.Run(ctx)
	if atomic.LoadInt32(&hits) != 0 {
		t.Fatalf("prometheus_only must not hit server, got %d requests", atomic.LoadInt32(&hits))
	}
	if sp.Len() != 0 {
		t.Fatalf("prometheus_only must not spool events, got %d", sp.Len())
	}
}

func TestEndToEndPushAndReplay(t *testing.T) {
	var (
		mu           sync.Mutex
		seenEventIDs = make(map[string]int)
		hbCount      int32
		fail         atomic.Bool
	)
	fail.Store(true)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/heartbeat":
			atomic.AddInt32(&hbCount, 1)
			w.WriteHeader(http.StatusAccepted)
			_, _ = io.WriteString(w, `{"accepted":true}`)
		case "/api/v1/events":
			if fail.Load() {
				http.Error(w, "down", http.StatusServiceUnavailable)
				return
			}
			var req client.EventBatchRequest
			_ = json.NewDecoder(r.Body).Decode(&req)
			mu.Lock()
			for _, e := range req.Events {
				seenEventIDs[e.EventID]++
			}
			mu.Unlock()
			w.WriteHeader(http.StatusAccepted)
			_ = json.NewEncoder(w).Encode(client.EventBatchResponse{Accepted: len(req.Events)})
		}
	}))
	defer srv.Close()

	dir := t.TempDir()
	cfg := &config.Config{
		AgentID:  "a1",
		Hostname: "h",
		Mode:     "hybrid",
		StateDir: dir,
		Server: config.ServerConfig{
			URL:               srv.URL,
			Token:             "t",
			HeartbeatInterval: config.Duration{Duration: 100 * time.Millisecond},
			BatchInterval:     config.Duration{Duration: 100 * time.Millisecond},
		},
		Metrics: config.MetricsConfig{Enabled: false},
		Checks: []config.CheckConfig{{
			ID:                "c1",
			Command:           []string{"sh", "-c", "echo hi"},
			Interval:          config.Duration{Duration: 150 * time.Millisecond},
			Timeout:           config.Duration{Duration: 100 * time.Millisecond},
			NotificationOwner: "server",
		}},
	}
	sp, err := spool.Open(filepath.Join(dir, "spool"), 100, 1<<20)
	if err != nil {
		t.Fatal(err)
	}
	a := &App{
		Cfg:     cfg,
		AgentID: "a1",
		Version: "test",
		Client:  client.New(cfg.Server.URL, "t", "a1", "test"),
		Spool:   sp,
		Metrics: metrics.New(),
		Log:     slog.New(slog.NewTextHandler(io.Discard, nil)),
	}

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { _ = a.Run(ctx); close(done) }()

	// Phase 1: server failing, events accumulate in spool.
	time.Sleep(400 * time.Millisecond)
	if sp.Len() == 0 {
		cancel()
		<-done
		t.Fatal("expected spooled events while server failing")
	}

	// Phase 2: server recovers.
	fail.Store(false)
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		mu.Lock()
		got := len(seenEventIDs)
		mu.Unlock()
		if got > 0 && sp.Len() == 0 {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}
	cancel()
	<-done

	mu.Lock()
	defer mu.Unlock()
	if len(seenEventIDs) == 0 {
		t.Fatal("no events delivered after recovery")
	}
	for id, n := range seenEventIDs {
		if n > 1 {
			t.Fatalf("event_id %s sent %d times (should be deduped by server, but agent should commit on first 2xx)", id, n)
		}
	}
	if atomic.LoadInt32(&hbCount) == 0 {
		t.Fatal("no heartbeats observed")
	}
}