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

package agent

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"google.golang.org/protobuf/proto"

	"git.opencomputing.cn/yumoqing/host-metrics-collector/pkg/collectorpb"
)

func newTestReporter(url, token string) *Reporter {
	r := NewReporter(url, token, "h-001", "0.1.0")
	// 测试中缩短重试退避与尝试次数，避免拖慢用例。
	r.maxAttempts = 3
	r.baseBackoff = time.Millisecond
	return r
}

func ackProto(t *testing.T, ok bool, seq int64, msg string) []byte {
	t.Helper()
	data, err := proto.Marshal(&collectorpb.Ack{Ok: ok, Seq: seq, Error: msg})
	if err != nil {
		t.Fatalf("marshal ack: %v", err)
	}
	return data
}

func TestReporterSendMetricsSuccess(t *testing.T) {
	var gotToken, gotHost, gotContentType string
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotToken = r.Header.Get(headerAgentToken)
		gotHost = r.Header.Get(headerHostID)
		gotContentType = r.Header.Get("Content-Type")
		w.Header().Set("Content-Type", contentTypeProtobuf)
		_, _ = w.Write(ackProto(t, true, 42, ""))
	}))
	defer ts.Close()

	r := newTestReporter(ts.URL, "test-token")
	ack, err := r.SendMetrics(context.Background(), &collectorpb.MetricBatch{HostId: "h-001", Seq: 42})
	if err != nil {
		t.Fatalf("SendMetrics: %v", err)
	}
	if !ack.Ok || ack.Seq != 42 {
		t.Fatalf("ack = %+v", ack)
	}
	if gotToken != "test-token" || gotHost != "h-001" || gotContentType != contentTypeProtobuf {
		t.Fatalf("headers: token=%q host=%q ct=%q", gotToken, gotHost, gotContentType)
	}
}

func TestReporterSendMetricsRetriesOnHTTPError(t *testing.T) {
	var calls int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if atomic.AddInt32(&calls, 1) < 3 {
			http.Error(w, "temporary", http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", contentTypeProtobuf)
		_, _ = w.Write(ackProto(t, true, 7, ""))
	}))
	defer ts.Close()

	r := newTestReporter(ts.URL, "test-token")
	ack, err := r.SendMetrics(context.Background(), &collectorpb.MetricBatch{HostId: "h-001", Seq: 7})
	if err != nil {
		t.Fatalf("SendMetrics: %v", err)
	}
	if !ack.Ok || ack.Seq != 7 {
		t.Fatalf("ack = %+v", ack)
	}
	if atomic.LoadInt32(&calls) != 3 {
		t.Fatalf("calls = %d, want 3", atomic.LoadInt32(&calls))
	}
}

func TestReporterSendMetricsAckNotOk(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", contentTypeProtobuf)
		_, _ = w.Write(ackProto(t, false, 0, "storage unavailable"))
	}))
	defer ts.Close()

	r := newTestReporter(ts.URL, "test-token")
	_, err := r.SendMetrics(context.Background(), &collectorpb.MetricBatch{HostId: "h-001", Seq: 1})
	if err == nil {
		t.Fatalf("expected error for ack not ok")
	}
	if !strings.Contains(err.Error(), "storage unavailable") {
		t.Fatalf("err = %v", err)
	}
}

func TestReporterSendMetricsExhaustsRetries(t *testing.T) {
	var calls int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&calls, 1)
		http.Error(w, "unavailable", http.StatusServiceUnavailable)
	}))
	defer ts.Close()

	r := newTestReporter(ts.URL, "test-token")
	_, err := r.SendMetrics(context.Background(), &collectorpb.MetricBatch{HostId: "h-001", Seq: 1})
	if err == nil {
		t.Fatalf("expected error after retries")
	}
	if atomic.LoadInt32(&calls) != 3 {
		t.Fatalf("calls = %d, want 3", atomic.LoadInt32(&calls))
	}
	if !strings.Contains(err.Error(), "failed after 3 attempts") {
		t.Fatalf("err = %v", err)
	}
}

func TestReporterSendLogsSuccess(t *testing.T) {
	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 log batch: %v", err)
		}
		if batch.HostId != "h-001" || batch.Seq != 9 || len(batch.Entries) != 1 {
			t.Errorf("batch = %+v", &batch)
		}
		w.Header().Set("Content-Type", contentTypeProtobuf)
		_, _ = w.Write(ackProto(t, true, 9, ""))
	}))
	defer ts.Close()

	r := newTestReporter(ts.URL, "test-token")
	ack, err := r.SendLogs(context.Background(), &collectorpb.LogBatch{
		HostId:  "h-001",
		Seq:     9,
		Entries: []*collectorpb.LogEntry{{Level: "ERROR", Message: "boom"}},
	})
	if err != nil {
		t.Fatalf("SendLogs: %v", err)
	}
	if !ack.Ok || ack.Seq != 9 {
		t.Fatalf("ack = %+v", ack)
	}
}

func TestReporterHeartbeatSuccess(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", contentTypeProtobuf)
		data, _ := proto.Marshal(&collectorpb.HeartbeatResponse{Ok: true, ServerTime: 1700000000})
		_, _ = w.Write(data)
	}))
	defer ts.Close()

	r := newTestReporter(ts.URL, "test-token")
	hb, err := r.Heartbeat(context.Background(), 0)
	if err != nil {
		t.Fatalf("Heartbeat: %v", err)
	}
	if !hb.Ok || hb.ServerTime != 1700000000 {
		t.Fatalf("hb = %+v", hb)
	}
}