restic-manager/internal/store/fleet_updates_test.go

package store

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/oklog/ulid/v2"
)

func ptrStr(s string) *string { return &s }

func seedFleetUser(t *testing.T, s *Store) string {
	t.Helper()
	id := ulid.Make().String()
	if err := s.CreateUser(context.Background(), User{
		ID: id, Username: "u-" + id[:6], PasswordHash: "x", Role: RoleAdmin,
	}); err != nil {
		t.Fatalf("create user: %v", err)
	}
	return id
}

func seedFleetHost(t *testing.T, s *Store, name string) string {
	t.Helper()
	id := ulid.Make().String()
	if err := s.CreateHost(context.Background(), Host{
		ID: id, Name: name, OS: "linux", Arch: "amd64",
		EnrolledAt: time.Now().UTC(),
	}, "tokenhash-"+id[:6], ""); err != nil {
		t.Fatalf("create host: %v", err)
	}
	return id
}

func TestCreateFleetUpdate_RefusesIfRunning(t *testing.T) {
	t.Parallel()
	s := openTestStore(t)
	uid := seedFleetUser(t, s)
	h1 := seedFleetHost(t, s, "h1")

	fu1 := FleetUpdate{ID: ulid.Make().String(), StartedByUserID: uid, TargetVersion: "v1"}
	if err := s.CreateFleetUpdate(context.Background(), fu1, []string{h1}); err != nil {
		t.Fatalf("create #1: %v", err)
	}
	fu2 := FleetUpdate{ID: ulid.Make().String(), StartedByUserID: uid, TargetVersion: "v2"}
	err := s.CreateFleetUpdate(context.Background(), fu2, []string{h1})
	if !errors.Is(err, ErrFleetUpdateRunning) {
		t.Fatalf("want ErrFleetUpdateRunning, got %v", err)
	}
}

func TestCreateFleetUpdate_HydrateRoundTrip(t *testing.T) {
	t.Parallel()
	s := openTestStore(t)
	uid := seedFleetUser(t, s)
	h1 := seedFleetHost(t, s, "h1")
	h2 := seedFleetHost(t, s, "h2")

	fu := FleetUpdate{ID: ulid.Make().String(), StartedByUserID: uid, TargetVersion: "v1.2.3"}
	if err := s.CreateFleetUpdate(context.Background(), fu, []string{h1, h2}); err != nil {
		t.Fatal(err)
	}

	got, hosts, err := s.GetFleetUpdate(context.Background(), fu.ID)
	if err != nil {
		t.Fatal(err)
	}
	if got.Status != "running" || got.TargetVersion != "v1.2.3" {
		t.Fatalf("parent: %+v", got)
	}
	if len(hosts) != 2 || hosts[0].Position != 0 || hosts[1].Position != 1 {
		t.Fatalf("hosts: %+v", hosts)
	}
	if hosts[0].Status != "pending" || hosts[1].Status != "pending" {
		t.Fatalf("hosts status: %+v", hosts)
	}
}

func TestSetFleetUpdateHostStatus_ProgressesAndStoresJobID(t *testing.T) {
	t.Parallel()
	s := openTestStore(t)
	uid := seedFleetUser(t, s)
	h := seedFleetHost(t, s, "h1")
	fu := FleetUpdate{ID: ulid.Make().String(), StartedByUserID: uid, TargetVersion: "v1"}
	_ = s.CreateFleetUpdate(context.Background(), fu, []string{h})

	jobID := ulid.Make().String()
	if err := s.CreateJob(context.Background(), Job{
		ID: jobID, HostID: h, Kind: "update",
		ActorKind: "user", ActorID: ptrStr(uid), CreatedAt: time.Now().UTC(),
	}); err != nil {
		t.Fatal(err)
	}

	if err := s.SetFleetUpdateHostStatus(context.Background(), fu.ID, h, "running", "", ""); err != nil {
		t.Fatal(err)
	}
	if err := s.SetFleetUpdateHostStatus(context.Background(), fu.ID, h, "succeeded", "", jobID); err != nil {
		t.Fatal(err)
	}
	_, hs, _ := s.GetFleetUpdate(context.Background(), fu.ID)
	if hs[0].Status != "succeeded" || hs[0].JobID != jobID {
		t.Fatalf("after succeed: %+v", hs[0])
	}

	pending, _ := s.ListPendingFleetUpdateHosts(context.Background(), fu.ID)
	if len(pending) != 0 {
		t.Fatalf("pending should be empty: %+v", pending)
	}
}

func TestHaltAndCompleteFleetUpdate(t *testing.T) {
	t.Parallel()
	s := openTestStore(t)
	uid := seedFleetUser(t, s)
	h := seedFleetHost(t, s, "h1")

	fu1 := FleetUpdate{ID: ulid.Make().String(), StartedByUserID: uid, TargetVersion: "v1"}
	_ = s.CreateFleetUpdate(context.Background(), fu1, []string{h})
	if err := s.HaltFleetUpdate(context.Background(), fu1.ID, "boom", time.Now().UTC()); err != nil {
		t.Fatal(err)
	}
	got, _, _ := s.GetFleetUpdate(context.Background(), fu1.ID)
	if got.Status != "halted" || got.HaltedReason != "boom" {
		t.Fatalf("after halt: %+v", got)
	}
	if got.CompletedAt == nil {
		t.Fatal("halted must stamp completed_at")
	}
	if active, _ := s.ActiveFleetUpdate(context.Background()); active != nil {
		t.Fatalf("halted should clear active: %+v", active)
	}

	// Now a fresh run can start.
	fu2 := FleetUpdate{ID: ulid.Make().String(), StartedByUserID: uid, TargetVersion: "v2"}
	if err := s.CreateFleetUpdate(context.Background(), fu2, []string{h}); err != nil {
		t.Fatalf("create after halt: %v", err)
	}
	if err := s.CompleteFleetUpdate(context.Background(), fu2.ID, time.Now().UTC()); err != nil {
		t.Fatal(err)
	}
	got, _, _ = s.GetFleetUpdate(context.Background(), fu2.ID)
	if got.Status != "completed" {
		t.Fatalf("after complete: %+v", got)
	}
}

func TestRunningUpdateJobForHost(t *testing.T) {
	t.Parallel()
	s := openTestStore(t)
	h := seedFleetHost(t, s, "h1")

	got, err := s.RunningUpdateJobForHost(context.Background(), h)
	if err != nil || got != "" {
		t.Fatalf("empty case: got=%q err=%v", got, err)
	}

	jobID := ulid.Make().String()
	if err := s.CreateJob(context.Background(), Job{
		ID: jobID, HostID: h, Kind: "update",
		ActorKind: "user", ActorID: ptrStr("u-1"), CreatedAt: time.Now().UTC(),
	}); err != nil {
		t.Fatal(err)
	}
	got, err = s.RunningUpdateJobForHost(context.Background(), h)
	if err != nil || got != jobID {
		t.Fatalf("queued case: got=%q err=%v", got, err)
	}

	// Mark succeeded → no longer "in flight".
	if err := s.MarkJobFinished(context.Background(), jobID, "succeeded", 0, nil, "", time.Now().UTC()); err != nil {
		t.Fatal(err)
	}
	got, err = s.RunningUpdateJobForHost(context.Background(), h)
	if err != nil || got != "" {
		t.Fatalf("after succeed: got=%q err=%v", got, err)
	}
}