felhom-agent/internal/localapi/controllerswap_test.go

package localapi

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"net/http/httptest"
	"regexp"
	"strings"
	"sync"
	"testing"
	"time"
)

func discardLogger() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) }

// fakeGuestExec simulates `pct exec` into a guest for the swap primitive. It models the image file,
// which images are present (pulled), which come up healthy, and the running container's image.
type fakeGuestExec struct {
	mu            sync.Mutex
	calls         [][]string
	imageFile     string          // /etc/felhom-controller-image content
	present       map[string]bool // images pulled into the guest
	good          map[string]bool // images that report healthy when running
	containerImg  string          // image the running container currently has
	failRestart   bool
	noHealthBlock bool // if set, .State.Health is absent ("none")
}

var writeImgRe = regexp.MustCompile(`'([^']+)' >`)

func (f *fakeGuestExec) GuestExec(_ context.Context, _ int, args ...string) (string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls = append(f.calls, args)
	switch {
	case len(args) >= 2 && args[0] == "cat" && args[1] == controllerImageFile:
		return f.imageFile + "\n", nil
	case len(args) >= 4 && args[0] == "docker" && args[1] == "image" && args[2] == "inspect":
		if f.present[args[3]] {
			return args[3], nil
		}
		return "", fmt.Errorf("no such image: %s", args[3])
	case len(args) >= 3 && args[0] == "bash" && args[1] == "-c":
		m := writeImgRe.FindStringSubmatch(args[2])
		if m == nil {
			return "", fmt.Errorf("fake: unparseable write cmd %q", args[2])
		}
		f.imageFile = m[1]
		return "", nil
	case len(args) >= 3 && args[0] == "systemctl" && args[1] == "restart":
		if f.failRestart {
			return "", fmt.Errorf("fake: systemctl restart failed")
		}
		f.containerImg = f.imageFile // bootstrap re-ran: container now runs the file's image
		return "", nil
	case len(args) >= 2 && args[0] == "docker" && args[1] == "inspect":
		if f.containerImg == "" {
			return "", fmt.Errorf("no such container")
		}
		health := "healthy"
		if f.noHealthBlock {
			health = "none"
		} else if !f.good[f.containerImg] {
			health = "unhealthy"
		}
		return fmt.Sprintf("true|%s|%s", health, f.containerImg), nil
	}
	return "", fmt.Errorf("fake: unexpected exec %v", args)
}

func (f *fakeGuestExec) wrote(image string) bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, c := range f.calls {
		if len(c) >= 3 && c[0] == "bash" && strings.Contains(c[2], "'"+image+"'") {
			return true
		}
	}
	return false
}

func newTestSwapper(t *testing.T, fe *fakeGuestExec) *ControllerSwapper {
	t.Helper()
	s := NewControllerSwapper(fe, t.TempDir(), nil)
	s.verifyTimeout = 200 * time.Millisecond
	s.verifyInterval = 5 * time.Millisecond
	return s
}

const (
	prevImg = "gitea.dooplex.hu/admin/felhom-controller:0.77.0"
	newImg  = "gitea.dooplex.hu/admin/felhom-controller:0.84.0"
)

func TestControllerSwap_Happy(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile: prevImg,
		present:   map[string]bool{newImg: true},
		good:      map[string]bool{newImg: true, prevImg: true},
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)

	if st.State != "done" {
		t.Fatalf("state = %q (err=%q), want done", st.State, st.Error)
	}
	if fe.imageFile != newImg {
		t.Errorf("image file = %q, want %q", fe.imageFile, newImg)
	}
	if st.Previous != prevImg {
		t.Errorf("previous = %q, want %q", st.Previous, prevImg)
	}
	if !fe.wrote(newImg) {
		t.Errorf("expected the new image to be written")
	}
	// state file persisted
	got, _ := s.LoadState(9201)
	if got == nil || got.State != "done" {
		t.Errorf("state file = %+v, want done", got)
	}
}

// The load-bearing failure path: an unhealthy target must roll back to the previous image.
func TestControllerSwap_RollbackOnUnhealthy(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile: prevImg,
		present:   map[string]bool{newImg: true},
		good:      map[string]bool{prevImg: true}, // newImg present but NEVER healthy
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)

	if st.State != "failed" {
		t.Fatalf("state = %q, want failed", st.State)
	}
	if fe.imageFile != prevImg {
		t.Errorf("image file = %q after rollback, want previous %q (guest left on bad image!)", fe.imageFile, prevImg)
	}
	if fe.containerImg != prevImg {
		t.Errorf("running container = %q after rollback, want %q", fe.containerImg, prevImg)
	}
	// restart called at least twice (swap + rollback)
	restarts := 0
	for _, c := range fe.calls {
		if len(c) >= 2 && c[0] == "systemctl" && c[1] == "restart" {
			restarts++
		}
	}
	if restarts < 2 {
		t.Errorf("systemctl restart called %d times, want ≥2 (swap + rollback)", restarts)
	}
}

func TestControllerSwap_ImageAbsent_NoSwap(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile: prevImg,
		present:   map[string]bool{}, // target NOT pulled
		good:      map[string]bool{prevImg: true},
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)

	if st.State != "failed" {
		t.Fatalf("state = %q, want failed", st.State)
	}
	if fe.imageFile != prevImg {
		t.Errorf("image file = %q, want unchanged %q (no swap on absent image)", fe.imageFile, prevImg)
	}
	if fe.wrote(newImg) {
		t.Errorf("must NOT write the image file when the target image is absent")
	}
}

func TestControllerSwap_HealthyWithNoHealthcheck(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile:     prevImg,
		present:       map[string]bool{newImg: true},
		noHealthBlock: true, // image defines no HEALTHCHECK → running is enough
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)
	if st.State != "done" {
		t.Fatalf("state = %q, want done (running + no healthcheck)", st.State)
	}
}

func TestControllerSwapHandler_BadImage(t *testing.T) {
	fe := &fakeGuestExec{present: map[string]bool{}}
	s := &Server{swap: newTestSwapper(t, fe), swapInFlight: map[int]bool{}, logger: discardLogger()}
	body := `{"image":"docker.io/evil/runme:latest"}`
	req := httptest.NewRequest("POST", "/controller/swap", strings.NewReader(body))
	rr := httptest.NewRecorder()
	s.handleControllerSwap(rr, req, 9201)
	if rr.Code != 400 {
		t.Fatalf("status = %d, want 400 for a non-controller image", rr.Code)
	}
	if len(fe.calls) != 0 {
		t.Errorf("no guest exec should run for a rejected image, got %v", fe.calls)
	}
}

func TestControllerSwapHandler_SingleFlight(t *testing.T) {
	fe := &fakeGuestExec{present: map[string]bool{}}
	s := &Server{swap: newTestSwapper(t, fe), swapInFlight: map[int]bool{9201: true}, logger: discardLogger()}
	req := httptest.NewRequest("POST", "/controller/swap", strings.NewReader(`{"image":"`+newImg+`"}`))
	rr := httptest.NewRecorder()
	s.handleControllerSwap(rr, req, 9201)
	if rr.Code != 409 {
		t.Fatalf("status = %d, want 409 (swap already in progress)", rr.Code)
	}
}