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) } }