host-metrics-collector/internal/agent/agent_test.go

package agent

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"
	"time"

	"google.golang.org/protobuf/proto"

	"git.opencomputing.cn/yumoqing/host-metrics-collector/internal/checkpoint"
	"git.opencomputing.cn/yumoqing/host-metrics-collector/internal/config"
	"git.opencomputing.cn/yumoqing/host-metrics-collector/internal/model"
	"git.opencomputing.cn/yumoqing/host-metrics-collector/internal/wal"
	"git.opencomputing.cn/yumoqing/host-metrics-collector/pkg/collectorpb"
)

func TestNextSeqMonotonic(t *testing.T) {
	a := &Agent{}
	seen := map[int64]bool{}
	for i := 0; i < 100; i++ {
		seq := a.nextMetricsSeq()
		if seen[seq] {
			t.Fatalf("duplicate metrics seq %d", seq)
		}
		seen[seq] = true
	}
	for i := 0; i < 100; i++ {
		seq := a.nextLogsSeq()
		if seq <= 0 {
			t.Fatalf("logs seq = %d", seq)
		}
	}
}

func TestInitSeqRestoresFromWAL(t *testing.T) {
	dir := t.TempDir()
	mw, err := wal.Open(filepath.Join(dir, "m.wal"))
	if err != nil {
		t.Fatalf("open metrics wal: %v", err)
	}
	defer mw.Close()
	lw, err := wal.Open(filepath.Join(dir, "l.wal"))
	if err != nil {
		t.Fatalf("open logs wal: %v", err)
	}
	defer lw.Close()

	// 模拟进程重启前已写入但尚未确认的批次。
	if err := mw.Append(3, []byte("m3")); err != nil {
		t.Fatalf("append m3: %v", err)
	}
	if err := mw.Append(5, []byte("m5")); err != nil {
		t.Fatalf("append m5: %v", err)
	}
	if err := lw.Append(2, []byte("l2")); err != nil {
		t.Fatalf("append l2: %v", err)
	}

	store := checkpoint.NewMemoryStore()
	// 已确认 seq 为 2，WAL 中还有 seq 3/5 未确认。
	if err := store.SetAckedSeq(metricsSeqKey+"h-001", 2); err != nil {
		t.Fatalf("set acked: %v", err)
	}

	a := &Agent{
		cfg:        config.AgentConfig{HostID: "h-001"},
		store:      store,
		metricsWAL: mw,
		logsWAL:    lw,
	}
	if err := a.initSeq(); err != nil {
		t.Fatalf("initSeq: %v", err)
	}
	if a.metricsSeq != 5 {
		t.Fatalf("metricsSeq = %d, want 5", a.metricsSeq)
	}
	if a.logsSeq != 2 {
		t.Fatalf("logsSeq = %d, want 2", a.logsSeq)
	}
}

// fakeGateway 返回一个接收 protobuf 上报并回 Ack 的 httptest 服务。
func fakeGateway(t *testing.T) (*httptest.Server, *int64, *[]string) {
	t.Helper()
	var ackedSeq int64
	var metrics []string
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get(headerAgentToken) != "test-token" {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		body, err := io.ReadAll(r.Body)
		if err != nil {
			http.Error(w, "read", http.StatusBadRequest)
			return
		}
		var batch collectorpb.MetricBatch
		if err := proto.Unmarshal(body, &batch); err != nil {
			http.Error(w, "bad proto", http.StatusBadRequest)
			return
		}
		for _, s := range batch.Samples {
			metrics = append(metrics, s.Name)
		}
		ackedSeq = batch.Seq
		w.Header().Set("Content-Type", contentTypeProtobuf)
		data, _ := proto.Marshal(&collectorpb.Ack{Ok: true, Seq: batch.Seq})
		_, _ = w.Write(data)
	}))
	return ts, &ackedSeq, &metrics
}

func TestFlushMetricsEndToEnd(t *testing.T) {
	ts, ackedSeq, metrics := fakeGateway(t)
	defer ts.Close()

	dir := t.TempDir()
	cfg := config.DefaultAgentConfig()
	cfg.GatewayURL = ts.URL
	cfg.AgentToken = "test-token"
	cfg.HostID = "h-001"
	cfg.CheckpointDriver = "memory"
	cfg.CheckpointFile = filepath.Join(dir, "cp")

	a, err := New(cfg)
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	defer a.Close()

	at := time.Now()
	samples := []model.Sample{
		model.NewSample(model.MetricCPUUsage, 42.5, at, map[string]string{model.LabelHostID: "h-001"}),
		model.NewSample(model.MetricMemUsedPercent, 66.6, at, map[string]string{model.LabelHostID: "h-001"}),
	}
	if err := a.flushMetrics(context.Background(), samples); err != nil {
		t.Fatalf("flushMetrics: %v", err)
	}

	if got := *ackedSeq; got != 1 {
		t.Fatalf("acked seq = %d, want 1", got)
	}
	if len(*metrics) != 2 {
		t.Fatalf("metrics = %v", *metrics)
	}
	if got, _ := a.store.GetAckedSeq(metricsSeqKey + "h-001"); got != 1 {
		t.Fatalf("local acked seq = %d, want 1", got)
	}
}

func TestFlushLogsEndToEnd(t *testing.T) {
	var gotSeq int64
	var gotMsgs []string
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var batch collectorpb.LogBatch
		if err := proto.Unmarshal(body, &batch); err != nil {
			t.Errorf("unmarshal: %v", err)
		}
		gotSeq = batch.Seq
		for _, e := range batch.Entries {
			gotMsgs = append(gotMsgs, e.Message)
		}
		w.Header().Set("Content-Type", contentTypeProtobuf)
		data, _ := proto.Marshal(&collectorpb.Ack{Ok: true, Seq: batch.Seq})
		_, _ = w.Write(data)
	}))
	defer ts.Close()

	dir := t.TempDir()
	cfg := config.DefaultAgentConfig()
	cfg.GatewayURL = ts.URL
	cfg.AgentToken = "test-token"
	cfg.HostID = "h-001"
	cfg.CheckpointDriver = "memory"
	cfg.CheckpointFile = filepath.Join(dir, "cp")

	a, err := New(cfg)
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	defer a.Close()

	records := []model.LogRecord{
		{Timestamp: time.Now().Unix(), Level: "ERROR", Message: "boom", File: "/var/log/app.log", Offset: 10},
	}
	if err := a.flushLogs(context.Background(), records); err != nil {
		t.Fatalf("flushLogs: %v", err)
	}
	if gotSeq != 1 || len(gotMsgs) != 1 || gotMsgs[0] != "boom" {
		t.Fatalf("seq=%d msgs=%v", gotSeq, gotMsgs)
	}
	if got, _ := a.store.GetAckedSeq(logsSeqKey + "h-001"); got != 1 {
		t.Fatalf("local acked seq = %d, want 1", got)
	}
}