host-metrics-collector/internal/gateway/server_test.go

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