Vetting/internal/api/agent_handlers_test.go

package api_test

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strconv"
	"testing"

	"github.com/go-chi/chi/v5"

	"vetting/internal/api"
	"vetting/internal/db"
	"vetting/internal/model"
	"vetting/internal/orchestrator"
	"vetting/internal/store"
)

func setupAgent(t *testing.T) (*api.Agent, int64, string) {
	t.Helper()
	path := filepath.Join(t.TempDir(), "vetting.db")
	conn, err := db.Open(path)
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	t.Cleanup(func() { _ = conn.Close() })

	hosts := &store.Hosts{DB: conn}
	runs := &store.Runs{DB: conn}
	meas := &store.Measurements{DB: conn}

	hostID, err := hosts.Create(context.Background(), model.Host{
		Name:             "t-host",
		MAC:              "aa:bb:cc:dd:ee:01",
		WoLBroadcastIP:   "10.0.0.255",
		WoLPort:          9,
		ExpectedSpecYAML: "memory:\n  total_gib: 16\n",
	})
	if err != nil {
		t.Fatalf("create host: %v", err)
	}
	plain, hash, err := orchestrator.IssueRunToken()
	if err != nil {
		t.Fatalf("issue token: %v", err)
	}
	runID, err := runs.Create(context.Background(), hostID, hash)
	if err != nil {
		t.Fatalf("create run: %v", err)
	}
	return &api.Agent{
		Hosts:        hosts,
		Runs:         runs,
		Measurements: meas,
	}, runID, plain
}

func routedRequest(runID int64, method, path string, body []byte) *http.Request {
	req := httptest.NewRequest(method, path, bytes.NewReader(body))
	// chi.URLParam is read from chi's context routing; fake that here.
	rctx := chi.NewRouteContext()
	rctx.URLParams.Add("id", strconv.FormatInt(runID, 10))
	return req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx))
}

func TestSensorPersistsBatch(t *testing.T) {
	a, runID, token := setupAgent(t)
	batch := api.SensorBatch{Samples: []api.SensorSample{
		{Kind: "thermal", Key: "cpu", Value: 47.5, Unit: "C"},
		{Kind: "iperf", Key: "throughput_mbps", Value: 938.2, Unit: "Mbps"},
	}}
	buf, _ := json.Marshal(batch)
	req := routedRequest(runID, http.MethodPost, "/api/v1/runs/"+strconv.FormatInt(runID, 10)+"/sensor", buf)
	req.Header.Set("Authorization", "Bearer "+token)
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()
	a.Sensor(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %d, body = %q", rr.Code, rr.Body.String())
	}
	rows, err := a.Measurements.ListForRun(context.Background(), runID)
	if err != nil {
		t.Fatalf("ListForRun: %v", err)
	}
	if len(rows) != 2 {
		t.Fatalf("expected 2 measurements, got %d", len(rows))
	}
}

func TestSensorRejectsBadToken(t *testing.T) {
	a, runID, _ := setupAgent(t)
	body, _ := json.Marshal(api.SensorBatch{})
	req := routedRequest(runID, http.MethodPost, "/api/v1/runs/"+strconv.FormatInt(runID, 10)+"/sensor", body)
	req.Header.Set("Authorization", "Bearer wrong-token")
	rr := httptest.NewRecorder()
	a.Sensor(rr, req)
	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want 401", rr.Code)
	}
}

// TestHeartbeatShutdownWhenCompleted: once the orchestrator has flipped
// the run into Completed, the next heartbeat response must carry
// cmd=shutdown so the agent powers the host down.
func TestHeartbeatShutdownWhenCompleted(t *testing.T) {
	a, runID, token := setupAgent(t)
	// Wire a runner so Heartbeat's TouchHeartbeat call doesn't nil-panic.
	a.Runner = &orchestrator.Runner{Runs: a.Runs, Hosts: a.Hosts, Stages: &store.Stages{DB: a.Runs.DB}}
	if err := a.Runs.SetState(context.Background(), runID, model.StateCompleted); err != nil {
		t.Fatalf("set state: %v", err)
	}
	req := routedRequest(runID, http.MethodPost, "/api/v1/runs/"+strconv.FormatInt(runID, 10)+"/heartbeat", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	rr := httptest.NewRecorder()
	a.Heartbeat(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %d, body = %s", rr.Code, rr.Body.String())
	}
	var resp map[string]any
	if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if resp["cmd"] != "shutdown" {
		t.Fatalf("cmd = %v, want shutdown", resp["cmd"])
	}
}