restic-manager/internal/server/http/run_group_bandwidth_test.go

// run_group_bandwidth_test.go — covers the per-job bandwidth override
// that operators can set via the Run-now form's "Limit bandwidth for
// this run" disclosure (P2R-13b).
package http

import (
	"context"
	"encoding/json"
	stdhttp "net/http"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"
	"github.com/oklog/ulid/v2"

	"gitea.dcglab.co.uk/steve/restic-manager/internal/api"
	"gitea.dcglab.co.uk/steve/restic-manager/internal/store"
)

// TestRunSourceGroupBandwidthOverride: connect a fake agent, POST the
// per-group Run-now endpoint with bandwidth_up_kbps=512, assert the
// dispatched command.run carries it.
func TestRunSourceGroupBandwidthOverride(t *testing.T) {
	t.Parallel()
	srv, ts, st := rawTestServer(t)
	hostID, token := enrolHostForWS(t, srv, st, "bw-host")

	// Pre-seed an init job so auto-init doesn't fire on hello and
	// pollute our envelope sequence.
	if err := st.CreateJob(context.Background(), store.Job{
		ID: ulid.Make().String(), HostID: hostID, Kind: "init",
		ActorKind: "system", CreatedAt: time.Now().UTC(),
	}); err != nil {
		t.Fatalf("seed init: %v", err)
	}

	gid := ulid.Make().String()
	if err := st.CreateSourceGroup(context.Background(), &store.SourceGroup{
		ID: gid, HostID: hostID, Name: "etc", Includes: []string{"/etc"},
	}); err != nil {
		t.Fatalf("group: %v", err)
	}

	c := agentDial(t, srv, ts, hostID, token)
	sendHello(t, c, "bw-host")
	// Drain on-hello burst before issuing the run-now.
	_ = drainUntil(t, c, api.MsgScheduleSet)

	cookie := loginAsAdmin(t, st)
	form := url.Values{
		"bandwidth_up_kbps":   {"512"},
		"bandwidth_down_kbps": {"256"},
	}
	req, _ := stdhttp.NewRequest("POST",
		ts.URL+"/hosts/"+hostID+"/source-groups/"+gid+"/run",
		strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")
	req.AddCookie(cookie)
	res, err := stdhttp.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	res.Body.Close()
	if res.StatusCode != stdhttp.StatusAccepted {
		t.Fatalf("status: got %d, want 202", res.StatusCode)
	}

	// Read the dispatched command.run; assert overrides are present.
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		ctx, cancel := context.WithTimeout(context.Background(), 800*time.Millisecond)
		mt, raw, rerr := c.Read(ctx)
		cancel()
		if rerr != nil {
			break
		}
		if mt != websocket.MessageText {
			continue
		}
		var env api.Envelope
		_ = json.Unmarshal(raw, &env)
		if env.Type != api.MsgCommandRun {
			continue
		}
		var p api.CommandRunPayload
		if err := env.UnmarshalPayload(&p); err != nil {
			t.Fatalf("unmarshal: %v", err)
		}
		if p.Kind != api.JobBackup {
			continue
		}
		if p.BandwidthUpKBps == nil || *p.BandwidthUpKBps != 512 {
			t.Fatalf("BandwidthUpKBps: got %v, want 512", p.BandwidthUpKBps)
		}
		if p.BandwidthDownKBps == nil || *p.BandwidthDownKBps != 256 {
			t.Fatalf("BandwidthDownKBps: got %v, want 256", p.BandwidthDownKBps)
		}
		return
	}
	t.Fatal("timed out waiting for command.run with bandwidth override")
}

// TestRunSourceGroupBandwidthRejectsNegative: invalid value → 400.
func TestRunSourceGroupBandwidthRejectsNegative(t *testing.T) {
	t.Parallel()
	_, url2, st := newTestServerWithHub(t)
	cookie := loginAsAdmin(t, st)
	hostID := makeHost(t, st, "bw-rej-host")
	gid := ulid.Make().String()
	if err := st.CreateSourceGroup(context.Background(), &store.SourceGroup{
		ID: gid, HostID: hostID, Name: "etc", Includes: []string{"/etc"},
	}); err != nil {
		t.Fatalf("group: %v", err)
	}
	form := url.Values{"bandwidth_up_kbps": {"-1"}}
	req, _ := stdhttp.NewRequest("POST",
		url2+"/hosts/"+hostID+"/source-groups/"+gid+"/run",
		strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")
	req.AddCookie(cookie)
	res, err := stdhttp.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	defer res.Body.Close()
	if res.StatusCode != stdhttp.StatusBadRequest {
		t.Fatalf("status: got %d, want 400", res.StatusCode)
	}
}