package gateway import ( "bytes" "encoding/json" "io" "net/http" "net/http/httptest" "strings" "testing" "time" "google.golang.org/protobuf/proto" "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/pkg/collectorpb" ) func newTestServer() (*httptest.Server, *Store) { cfg := config.GatewayConfig{Addr: "127.0.0.1:4318", AgentToken: "test-token", MaxMetricsPerHost: 100, MaxLogsPerHost: 100} store := NewStore(cfg.MaxMetricsPerHost, cfg.MaxLogsPerHost) srv := httptest.NewServer(NewServer(cfg, store).Handler()) return srv, store } func agentRequest(t *testing.T, method, url string, body []byte, token string) *http.Response { t.Helper() req, err := http.NewRequest(method, url, bytes.NewReader(body)) if err != nil { t.Fatalf("new request: %v", err) } if token != "" { req.Header.Set(headerAgentToken, token) } if body != nil { req.Header.Set("Content-Type", contentTypeProtobuf) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("do request: %v", err) } return resp } func readAll(t *testing.T, resp *http.Response) []byte { t.Helper() defer resp.Body.Close() data, err := io.ReadAll(resp.Body) if err != nil { t.Fatalf("read body: %v", err) } return data } func decodeJSON(t *testing.T, resp *http.Response, out any) { t.Helper() defer resp.Body.Close() if err := json.NewDecoder(resp.Body).Decode(out); err != nil { t.Fatalf("decode json: %v", err) } } func TestAgentMetricsRoundTrip(t *testing.T) { ts, _ := newTestServer() defer ts.Close() batch := &collectorpb.MetricBatch{ HostId: "h-001", Seq: 42, SentAt: time.Now().Unix(), Samples: []*collectorpb.MetricSample{ {Name: "cpu.usage", Value: 55.5, Timestamp: time.Now().Unix(), Labels: map[string]string{"host_id": "h-001"}}, }, } body, _ := proto.Marshal(batch) resp := agentRequest(t, http.MethodPost, ts.URL+"/v1/agent/metrics", body, "test-token") if resp.StatusCode != http.StatusOK { t.Fatalf("status = %d, want 200", resp.StatusCode) } var ack collectorpb.Ack if err := proto.Unmarshal(readAll(t, resp), &ack); err != nil { t.Fatalf("unmarshal ack: %v", err) } if !ack.Ok || ack.Seq != 42 { t.Fatalf("ack ok=%v seq=%d, want ok=true seq=42", ack.Ok, ack.Seq) } } func TestAgentMetricsAuth(t *testing.T) { ts, _ := newTestServer() defer ts.Close() body, _ := proto.Marshal(&collectorpb.MetricBatch{HostId: "h-001", Seq: 1}) resp := agentRequest(t, http.MethodPost, ts.URL+"/v1/agent/metrics", body, "wrong-token") if resp.StatusCode != http.StatusUnauthorized { t.Fatalf("status = %d, want 401", resp.StatusCode) } } func TestAgentMetricsStoresSamples(t *testing.T) { ts, store := newTestServer() defer ts.Close() at := time.Now().Unix() batch := &collectorpb.MetricBatch{ HostId: "h-001", Seq: 7, SentAt: at, Samples: []*collectorpb.MetricSample{ {Name: "cpu.usage", Value: 90, Timestamp: at, Labels: map[string]string{model.LabelHostID: "h-001", model.LabelMetricGroup: model.GroupCore}}, }, } body, _ := proto.Marshal(batch) resp := agentRequest(t, http.MethodPost, ts.URL+"/v1/agent/metrics", body, "test-token") if resp.StatusCode != http.StatusOK { t.Fatalf("status = %d", resp.StatusCode) } readAll(t, resp) series, err := store.HostMetrics("h-001", "cpu.usage", time.Time{}, time.Time{}) if err != nil { t.Fatalf("HostMetrics: %v", err) } if len(series) != 1 || len(series[0].Values) != 1 || series[0].Values[0].Value != 90 { t.Fatalf("series = %+v", series) } if got := store.LastAckedSeq("h-001"); got != 7 { t.Fatalf("acked seq = %d, want 7", got) } } func TestAgentLogsStoresEntries(t *testing.T) { ts, store := newTestServer() defer ts.Close() batch := &collectorpb.LogBatch{ HostId: "h-001", Seq: 3, Entries: []*collectorpb.LogEntry{ {Timestamp: time.Now().Unix(), Level: "ERROR", Source: "/var/log/app.log", Message: "boom", File: "/var/log/app.log", Offset: 100}, }, } body, _ := proto.Marshal(batch) resp := agentRequest(t, http.MethodPost, ts.URL+"/v1/agent/logs", body, "test-token") if resp.StatusCode != http.StatusOK { t.Fatalf("status = %d", resp.StatusCode) } readAll(t, resp) total, items, err := store.HostLogs("h-001", "ERROR", "boom", time.Time{}, time.Time{}, 1, 20) if err != nil { t.Fatalf("HostLogs: %v", err) } if total != 1 || len(items) != 1 || items[0].Message != "boom" { t.Fatalf("total=%d items=%+v", total, items) } } func TestHeartbeat(t *testing.T) { ts, _ := newTestServer() defer ts.Close() body, _ := proto.Marshal(&collectorpb.HeartbeatRequest{HostId: "h-001", AgentVersion: "0.1.0"}) resp := agentRequest(t, http.MethodPost, ts.URL+"/v1/agent/heartbeat", body, "test-token") if resp.StatusCode != http.StatusOK { t.Fatalf("status = %d", resp.StatusCode) } var hb collectorpb.HeartbeatResponse if err := proto.Unmarshal(readAll(t, resp), &hb); err != nil { t.Fatalf("unmarshal heartbeat: %v", err) } if !hb.Ok { t.Fatalf("heartbeat not ok") } } func TestCheckpoint(t *testing.T) { ts, store := newTestServer() defer ts.Close() store.SetLastAckedSeq("h-001", 99) resp := agentRequest(t, http.MethodGet, ts.URL+"/v1/agent/checkpoint/h-001", nil, "test-token") var cp collectorpb.CheckpointResponse if err := proto.Unmarshal(readAll(t, resp), &cp); err != nil { t.Fatalf("unmarshal checkpoint: %v", err) } if cp.LastAckedSeq != 99 { t.Fatalf("checkpoint = %d, want 99", cp.LastAckedSeq) } } func TestRESTHostMetrics(t *testing.T) { ts, store := newTestServer() defer ts.Close() at := time.Now().Unix() store.AddMetrics("h-001", 1, []model.Sample{ {Name: "cpu.usage", Value: 12.5, Timestamp: at, Labels: map[string]string{model.LabelHostID: "h-001"}}, }) resp, err := http.Get(ts.URL + "/api/v1/hosts/h-001/metrics?metric=cpu.usage") if err != nil { t.Fatalf("get: %v", err) } var envelope struct { Code int `json:"code"` Message string `json:"message"` Data []HostMetricSeries `json:"data"` } decodeJSON(t, resp, &envelope) if envelope.Code != 0 || len(envelope.Data) != 1 || envelope.Data[0].Values[0].Value != 12.5 { t.Fatalf("envelope = %+v", envelope) } } func TestRESTHostLogsPagination(t *testing.T) { ts, store := newTestServer() defer ts.Close() for i := 0; i < 25; i++ { store.AddLogs("h-001", int64(i+1), []model.LogRecord{{Timestamp: time.Now().Unix(), Level: "INFO", Message: "line"}}) } resp, err := http.Get(ts.URL + "/api/v1/hosts/h-001/logs?page=2&page_size=10") if err != nil { t.Fatalf("get: %v", err) } var envelope struct { Code int `json:"code"` Data struct { Total int `json:"total"` Items []LogEntryDTO `json:"items"` } `json:"data"` } decodeJSON(t, resp, &envelope) if envelope.Code != 0 || envelope.Data.Total != 25 || len(envelope.Data.Items) != 10 { t.Fatalf("envelope = %+v", envelope) } } func TestRESTMetricsQueryInstant(t *testing.T) { ts, store := newTestServer() defer ts.Close() at := time.Now().Unix() store.AddMetrics("h-001", 1, []model.Sample{ {Name: "mem.used_percent", Value: 88, Timestamp: at, Labels: map[string]string{model.LabelHostID: "h-001", "service": "web"}}, }) resp, err := http.Get(ts.URL + "/api/v1/metrics/query/instant?metric=mem.used_percent&labels=service=web") if err != nil { t.Fatalf("get: %v", err) } var envelope struct { Code int `json:"code"` Data struct { Result []MetricPoint `json:"result"` } `json:"data"` } decodeJSON(t, resp, &envelope) if envelope.Code != 0 || len(envelope.Data.Result) != 1 || envelope.Data.Result[0].Value != 88 { t.Fatalf("envelope = %+v", envelope) } } func TestRESTLogsSearch(t *testing.T) { ts, store := newTestServer() defer ts.Close() store.AddLogs("h-001", 1, []model.LogRecord{{Timestamp: time.Now().Unix(), Level: "ERROR", Message: "disk full"}}) store.AddLogs("h-002", 1, []model.LogRecord{{Timestamp: time.Now().Unix(), Level: "INFO", Message: "ok"}}) body := `{"host_ids":["h-001"],"level":"ERROR","keyword":"disk"}` resp, err := http.Post(ts.URL+"/api/v1/logs/search", "application/json", strings.NewReader(body)) if err != nil { t.Fatalf("post: %v", err) } var envelope struct { Code int `json:"code"` Data struct { Total int `json:"total"` Items []LogEntryDTO `json:"items"` } `json:"data"` } decodeJSON(t, resp, &envelope) if envelope.Code != 0 || envelope.Data.Total != 1 || envelope.Data.Items[0].Message != "disk full" { t.Fatalf("envelope = %+v", envelope) } } func TestParseTimeAndLabels(t *testing.T) { if _, err := parseTime("1700000000"); err != nil { t.Fatalf("unix parse: %v", err) } if _, err := parseTime("2025-01-01T00:00:00Z"); err != nil { t.Fatalf("rfc3339 parse: %v", err) } if _, err := parseTime("bad"); err == nil { t.Fatalf("expected error for bad time") } labels, err := parseLabels("host_id=h-001,service=web") if err != nil { t.Fatalf("parseLabels: %v", err) } if labels["service"] != "web" { t.Fatalf("labels = %v", labels) } if _, err := parseLabels("bad"); err == nil { t.Fatalf("expected error for bad labels") } }