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