feat(tests): missing vm tests (#196)

go.mod

@@ -22,6 +22,7 @@ require (
 	github.com/spf13/cobra v1.10.2
 	github.com/spf13/viper v1.21.0
 	github.com/stretchr/testify v1.11.1
+	go.uber.org/goleak v1.3.0
 	golang.org/x/crypto v0.50.0
 	gopkg.in/ini.v1 v1.67.2
 )

spinifex/vm/crash_recovery_test.go

@@ -391,3 +391,54 @@ func TestRestartCrashedInstance_AllocateFailure(t *testing.T) {
 	assert.Equal(t, 1, instance.Health.RestartCount,
 		"RestartCount is incremented before the allocate attempt")
 }
+
+// TestRestartCrashedInstance_RunFailureRollback covers the post-Allocate
+// Run-failure path: Allocate succeeds, transition to Pending succeeds,
+// then m.Run fails (Mount error). The rollback must deallocate the
+// reservation and transition the instance back to StateError. A
+// regression that skipped the deallocate would leak one reservation
+// per crash-then-restart-fail cycle until the daemon restarts.
+func TestRestartCrashedInstance_RunFailureRollback(t *testing.T) {
+	t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
+	m, rc, rt, _ := crashTestManager(t)
+
+	// Wire a VolumeMounter that fails so m.Run returns an error after the
+	// successful Allocate + Pending transition. SetDeps replaces deps
+	// entirely, so re-supply the existing wiring from crashTestManager.
+	mounter := &fakeVolumeMounter{mountErr: errors.New("mount failed during restart")}
+	m.SetDeps(Deps{
+		NodeID:          "test-node",
+		Resources:       rc,
+		VolumeMounter:   mounter,
+		InstanceTypes:   fakeInstanceTypeResolver{"t3.micro": {VCPUs: 1, MemoryMiB: 1024, Architecture: "x86_64"}},
+		TransitionState: rt.apply,
+	})
+
+	// Pre-allocate one reservation so the deallocate-on-rollback is
+	// observable as a return-to-baseline rather than a negative count.
+	require.NoError(t, rc.Allocate("t3.micro"))
+	baseline := rc.allocateCount("t3.micro")
+
+	instance := &VM{
+		ID:           "i-rollback",
+		Status:       StateError,
+		InstanceType: "t3.micro",
+	}
+	m.Insert(instance)
+
+	m.RestartCrashedInstance(instance)
+
+	assert.Equal(t, baseline, rc.allocateCount("t3.micro"),
+		"Run-failure rollback must net allocations back to baseline")
+	assert.Equal(t, 1, rc.deallocateCount("t3.micro"),
+		"Deallocate must be invoked exactly once on Run-failure rollback")
+	assert.Equal(t, []string{"i-rollback"}, mounter.mounted,
+		"Mount must have been attempted (proves the launch progressed past the Pending transition)")
+
+	targets := rt.targets("i-rollback")
+	require.Len(t, targets, 2, "expected two transitions: Pending then back to Error")
+	assert.Equal(t, StatePending, targets[0])
+	assert.Equal(t, StateError, targets[1])
+	assert.Equal(t, StateError, m.Status(instance),
+		"instance must end in StateError after Run-failure rollback")
+}

spinifex/vm/lifecycle_test.go

@@ -1,12 +1,17 @@
 package vm
 
 import (
+	"errors"
 	"fmt"
+	"net"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 )
 
 func TestBuildBaseVMConfig(t *testing.T) {
@@ -226,3 +231,261 @@ func TestStartReturnsErrorWhenInstanceUnknown(t *testing.T) {
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "i-missing")
 }
+
+// launchTestManager wires the minimum deps needed to drive Manager.launch
+// far enough to exercise the early-exit paths and the VolumeMounter call
+// site. Returns the manager, the mounter (for assertions and onMount
+// hooks), and a counter of OnInstanceUp invocations.
+//
+// The success-path "OnInstanceUp fires exactly once" assertion is not
+// driven from a vm-package unit test — the launch happy path requires a
+// real QEMU process plus a working QMP socket (startQEMU + AttachQMP),
+// and the heartbeat goroutine spawned by AttachQMP blocks on a 30s sleep
+// with no cancellation seam, so a hermetic test would either leak the
+// heartbeat or stall for 30s. The "exactly once" contract is covered by
+// the daemon-side e2e tests that drive the full RunInstances pipeline.
+// The counter is retained so the negative direction (no fire on aborted
+// or failed launches) can be asserted alongside the existing tests.
+func launchTestManager(t *testing.T) (*Manager, *fakeVolumeMounter, *int) {
+	t.Helper()
+	t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
+	mounter := &fakeVolumeMounter{}
+	upCalls := 0
+	m := NewManager()
+	m.SetDeps(Deps{
+		VolumeMounter: mounter,
+		Hooks: ManagerHooks{
+			OnInstanceUp: func(*VM) error { upCalls++; return nil },
+		},
+	})
+	return m, mounter, &upCalls
+}
+
+// TestRun_LaunchStillValid_FirstCheck covers the first launchStillValid
+// race-check (lifecycle.go:77): when a concurrent terminate has flipped
+// status out of {Pending, Stopped, Provisioning} before Run starts, Run
+// must return nil immediately without mounting volumes or firing the
+// OnInstanceUp hook — the terminate goroutine owns cleanup.
+func TestRun_LaunchStillValid_FirstCheck(t *testing.T) {
+	for _, status := range []InstanceState{
+		StateRunning, StateStopping, StateShuttingDown, StateTerminated, StateError,
+	} {
+		t.Run(string(status), func(t *testing.T) {
+			m, mounter, upCalls := launchTestManager(t)
+			instance := &VM{ID: "i-" + string(status), Status: status}
+			m.Insert(instance)
+
+			err := m.Run(instance)
+
+			require.NoError(t, err)
+			assert.Empty(t, mounter.mounted, "Mount must not be called when launch is aborted by first race check")
+			assert.Equal(t, 0, *upCalls, "OnInstanceUp must not fire when launch is aborted")
+			assert.Equal(t, status, m.Status(instance), "status must not change")
+		})
+	}
+}
+
+// TestRun_LaunchStillValid_SecondCheck covers the second launchStillValid
+// race-check (lifecycle.go:102): a terminate races in during Mount (which
+// can take 30+s on cold AMIs) and flips status to ShuttingDown. Run must
+// return nil after Mount completes, without proceeding to startQEMU or
+// firing OnInstanceUp.
+func TestRun_LaunchStillValid_SecondCheck(t *testing.T) {
+	m, mounter, upCalls := launchTestManager(t)
+	instance := &VM{ID: "i-flip", Status: StatePending}
+	m.Insert(instance)
+
+	mounter.onMount = func(v *VM) {
+		m.UpdateState(v.ID, func(vv *VM) { vv.Status = StateShuttingDown })
+	}
+
+	err := m.Run(instance)
+
+	require.NoError(t, err, "Run must return nil when concurrent terminate flips status during Mount")
+	assert.Equal(t, []string{"i-flip"}, mounter.mounted, "Mount must run exactly once before the second race check")
+	assert.Equal(t, 0, *upCalls, "OnInstanceUp must not fire when the second race check aborts launch")
+	assert.Equal(t, StateShuttingDown, m.Status(instance))
+}
+
+// TestRun_VolumeMounterError_Propagates covers the Mount-failure branch
+// (lifecycle.go:93): the error must bubble up unchanged, no startQEMU, no
+// hook fires, and the manager state is unaffected.
+func TestRun_VolumeMounterError_Propagates(t *testing.T) {
+	m, mounter, upCalls := launchTestManager(t)
+	sentinel := errors.New("mount failed")
+	mounter.mountErr = sentinel
+
+	instance := &VM{ID: "i-mount-fail", Status: StatePending}
+	m.Insert(instance)
+
+	err := m.Run(instance)
+
+	require.ErrorIs(t, err, sentinel)
+	assert.Equal(t, []string{"i-mount-fail"}, mounter.mounted)
+	assert.Equal(t, 0, *upCalls)
+	assert.Equal(t, StatePending, m.Status(instance), "status must be unchanged on Mount failure")
+}
+
+// TestRun_AlreadyRunningPID_ReturnsError covers the live-PID guard
+// (lifecycle.go:81-91): when the PID file refers to a live process, Run
+// must return an error before touching volumes or firing hooks.
+func TestRun_AlreadyRunningPID_ReturnsError(t *testing.T) {
+	m, mounter, upCalls := launchTestManager(t)
+
+	instance := &VM{ID: "i-live-pid", Status: StatePending}
+	m.Insert(instance)
+
+	pidDir := os.Getenv("XDG_RUNTIME_DIR")
+	pidFile := filepath.Join(pidDir, instance.ID+".pid")
+	require.NoError(t, os.WriteFile(pidFile, fmt.Appendf(nil, "%d", os.Getpid()), 0o600))
+
+	err := m.Run(instance)
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "already running")
+	assert.Empty(t, mounter.mounted, "Mount must not run when an existing live PID is detected")
+	assert.Equal(t, 0, *upCalls)
+}
+
+// TestStart_DispatchesThroughLaunch asserts Start (the found-instance
+// branch) actually drives launch — not merely returns nil. A regression
+// that made Start a no-op after the lookup would still pass
+// TestStartReturnsErrorWhenInstanceUnknown; this test catches it by
+// requiring the launch-internal Mount call.
+func TestStart_DispatchesThroughLaunch(t *testing.T) {
+	m, mounter, _ := launchTestManager(t)
+	sentinel := errors.New("mount failed via Start")
+	mounter.mountErr = sentinel
+
+	instance := &VM{ID: "i-start-dispatch", Status: StateStopped}
+	m.Insert(instance)
+
+	err := m.Start(instance.ID)
+
+	require.ErrorIs(t, err, sentinel, "Start must dispatch through launch and surface Mount errors")
+	assert.Equal(t, []string{"i-start-dispatch"}, mounter.mounted)
+}
+
+// TestStart_AbortedByConcurrentTerminate is the Start-side mirror of the
+// first launchStillValid check: a status set outside the launchable
+// allowlist must abort the launch with no error and no Mount call.
+func TestStart_AbortedByConcurrentTerminate(t *testing.T) {
+	m, mounter, upCalls := launchTestManager(t)
+	instance := &VM{ID: "i-start-abort", Status: StateShuttingDown}
+	m.Insert(instance)
+
+	err := m.Start(instance.ID)
+
+	require.NoError(t, err)
+	assert.Empty(t, mounter.mounted)
+	assert.Equal(t, 0, *upCalls)
+}
+
+// TestLaunchStillValid_Allowlist locks in the launchStillValid allowlist
+// across every defined InstanceState plus the zero value. The parent
+// plan flagged: "a regression in the allowlist would silently break the
+// start-stopped path" — e.g. dropping StateStopped would make
+// StartInstances on a stopped VM no-op without surfacing an error.
+func TestLaunchStillValid_Allowlist(t *testing.T) {
+	tests := []struct {
+		status InstanceState
+		want   bool
+		why    string
+	}{
+		{StateProvisioning, true, "RunInstances entry: launches before status flip to Pending"},
+		{StatePending, true, "RunInstances/restore entry: launch must proceed"},
+		{StateStopped, true, "StartInstances re-launches a stopped VM through the same pipeline"},
+		{StateRunning, false, "already past launch — second Run is a no-op"},
+		{StateStopping, false, "stop in flight: stopCleanup owns teardown"},
+		{StateShuttingDown, false, "terminate in flight: terminateCleanup owns teardown"},
+		{StateTerminated, false, "terminal"},
+		{StateError, false, "RestartCrashedInstance drives transitions, not launch"},
+		{InstanceState(""), false, "zero value must reject"},
+	}
+
+	m := NewManager()
+	for _, tt := range tests {
+		t.Run(string(tt.status), func(t *testing.T) {
+			instance := &VM{ID: "i-" + string(tt.status), Status: tt.status}
+			got := m.launchStillValid(instance)
+			assert.Equal(t, tt.want, got, tt.why)
+		})
+	}
+}
+
+// startBrokenQMPListener accepts on a unix socket and immediately closes
+// every connection without sending a greeting, so qmp.NewQMPClient's
+// greeting decode fails. Returns the socket path and a stop function the
+// caller must invoke before goleak.VerifyNone so the accept goroutine
+// drains.
+func startBrokenQMPListener(t *testing.T) (string, func()) {
+	t.Helper()
+	sockPath := filepath.Join(t.TempDir(), "qmp.sock")
+	ln, err := net.Listen("unix", sockPath)
+	require.NoError(t, err)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			_ = conn.Close()
+		}
+	}()
+
+	stopped := false
+	stop := func() {
+		if stopped {
+			return
+		}
+		stopped = true
+		_ = ln.Close()
+		<-done
+	}
+	t.Cleanup(stop)
+	return sockPath, stop
+}
+
+// TestAttachQMP_DialFailure_NoHeartbeatLeak covers the regression the
+// parent plan flagged: a swap that started qmpHeartbeat before checking
+// the factory error would leak one goroutine per failed reconnect. With
+// no listener at the configured socket, net.Dial fails inside
+// newQMPClientWithHandshake and AttachQMP must return without spawning
+// the heartbeat or mutating instance.QMPClient.
+func TestAttachQMP_DialFailure_NoHeartbeatLeak(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	m := NewManager()
+	instance := &VM{
+		ID:     "i-dial-fail",
+		Config: Config{QMPSocket: filepath.Join(t.TempDir(), "missing.sock")},
+	}
+
+	err := m.AttachQMP(instance)
+
+	require.Error(t, err)
+	assert.Nil(t, instance.QMPClient, "QMPClient must remain unset on dial failure")
+}
+
+// TestAttachQMP_HandshakeFailure_NoHeartbeatLeak covers the path where
+// the unix socket exists but qmp.NewQMPClient's greeting decode fails
+// (server closed). The dial succeeded so the connection must be closed
+// before AttachQMP returns, the heartbeat must not start, and
+// instance.QMPClient must remain nil.
+func TestAttachQMP_HandshakeFailure_NoHeartbeatLeak(t *testing.T) {
+	sockPath, stopListener := startBrokenQMPListener(t)
+
+	m := NewManager()
+	instance := &VM{ID: "i-handshake-fail", Config: Config{QMPSocket: sockPath}}
+
+	err := m.AttachQMP(instance)
+
+	require.Error(t, err)
+	assert.Nil(t, instance.QMPClient, "QMPClient must remain unset on handshake failure")
+
+	stopListener()
+	goleak.VerifyNone(t)
+}

spinifex/vm/manager_deps_test.go

@@ -146,13 +146,21 @@ type fakeVolumeMounter struct {
 	mountErr                 error
 	mountOneErr              error
 	mountOneURI              string
+	// onMount fires synchronously inside Mount before the configured
+	// mountErr is returned. Used by lifecycle tests to simulate a
+	// concurrent terminate flipping VM.Status while Mount is in flight.
+	onMount func(*VM)
 }
 
 func (f *fakeVolumeMounter) Mount(v *VM) error {
 	f.mu.Lock()
 	f.mounted = append(f.mounted, v.ID)
 	err := f.mountErr
+	hook := f.onMount
 	f.mu.Unlock()
+	if hook != nil {
+		hook(v)
+	}
 	return err
 }
 

spinifex/vm/restore_test.go

@@ -0,0 +1,351 @@
+package vm
+
+import (
+	"errors"
+	"maps"
+	"net"
+	"path/filepath"
+	"sync/atomic"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/mulgadc/spinifex/spinifex/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// withNBDRequests returns a VM whose EBSRequests slice carries the given
+// URIs in the order supplied. Each request is named "vol-<i>" so test
+// failures can identify which entry tripped a probe.
+func withNBDRequests(id string, status InstanceState, uris ...string) *VM {
+	v := &VM{ID: id, Status: status}
+	v.EBSRequests.Requests = make([]types.EBSRequest, len(uris))
+	for i, uri := range uris {
+		v.EBSRequests.Requests[i].Name = "vol-" + string(rune('a'+i))
+		v.EBSRequests.Requests[i].NBDURI = uri
+	}
+	return v
+}
+
+func TestAreVolumeSocketsValid_NoRequests(t *testing.T) {
+	assert.True(t, AreVolumeSocketsValid(&VM{ID: "i-empty"}),
+		"a VM with zero NBD requests has nothing to probe and must report valid")
+}
+
+func TestAreVolumeSocketsValid_EmptyURI(t *testing.T) {
+	v := withNBDRequests("i-empty-uri", StateRunning, "")
+	assert.True(t, AreVolumeSocketsValid(v),
+		"empty NBDURI is skipped (volume not yet mounted), not probed")
+}
+
+func TestAreVolumeSocketsValid_TCPURI(t *testing.T) {
+	v := withNBDRequests("i-tcp", StateRunning, "nbd://10.0.0.1:10809")
+	assert.True(t, AreVolumeSocketsValid(v),
+		"TCP NBD URI is treated as valid — recovery path can't probe a remote viperblockd")
+}
+
+func TestAreVolumeSocketsValid_UnparseableURI(t *testing.T) {
+	v := withNBDRequests("i-junk", StateRunning, "not-an-nbd-uri")
+	assert.True(t, AreVolumeSocketsValid(v),
+		"unparseable URI must not block reconnect — fall through to the launch path's stricter check")
+}
+
+func TestAreVolumeSocketsValid_UnixSocketMissing(t *testing.T) {
+	missing := filepath.Join(t.TempDir(), "missing.sock")
+	v := withNBDRequests("i-no-sock", StateRunning, "nbd:unix:"+missing)
+	assert.False(t, AreVolumeSocketsValid(v),
+		"a unix NBD socket that does not accept connections is the orphan-QEMU signal")
+}
+
+func TestAreVolumeSocketsValid_UnixSocketLive(t *testing.T) {
+	sockPath := filepath.Join(t.TempDir(), "live.sock")
+	ln, err := net.Listen("unix", sockPath)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = ln.Close() })
+
+	v := withNBDRequests("i-live", StateRunning, "nbd:unix:"+sockPath)
+	assert.True(t, AreVolumeSocketsValid(v),
+		"a live unix listener at the configured path means viperblock survived the daemon restart")
+}
+
+func TestAreVolumeSocketsValid_OneOfManyMissing(t *testing.T) {
+	live := filepath.Join(t.TempDir(), "live.sock")
+	dead := filepath.Join(t.TempDir(), "dead.sock")
+	ln, err := net.Listen("unix", live)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = ln.Close() })
+
+	v := withNBDRequests("i-mixed", StateRunning, "nbd:unix:"+live, "nbd:unix:"+dead)
+	assert.False(t, AreVolumeSocketsValid(v),
+		"a single dead socket is enough to invalidate the whole instance")
+}
+
+// classifyTestManager wires a manager with an in-memory StateStore plus
+// the InstanceTypes resolver and Resources controller that
+// classifyRestoredInstances consults. Returns the manager and the
+// fakes for assertion.
+func classifyTestManager(t *testing.T) (*Manager, *fakeStateStore, *fakeResourceController) {
+	t.Helper()
+	t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
+	store := newFakeStateStore()
+	rc := newFakeResourceController()
+	m := NewManager()
+	m.SetDeps(Deps{
+		NodeID:        "test-node",
+		StateStore:    store,
+		Resources:     rc,
+		InstanceTypes: fakeInstanceTypeResolver{"t3.micro": {VCPUs: 1, MemoryMiB: 1024, Architecture: "x86_64"}},
+	})
+	return m, store, rc
+}
+
+func TestClassifyRestoredInstances_TerminatedMigrates(t *testing.T) {
+	m, store, _ := classifyTestManager(t)
+	v := &VM{ID: "i-term", Status: StateTerminated, InstanceType: "t3.micro"}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	assert.Empty(t, toLaunch, "terminated must not be queued for relaunch")
+	assert.NotNil(t, store.terminated[v.ID], "terminated must be migrated to the shared bucket")
+	_, stillLocal := m.Get(v.ID)
+	assert.False(t, stillLocal, "terminated must be removed from the local map")
+}
+
+func TestClassifyRestoredInstances_StoppedMigrates(t *testing.T) {
+	m, store, _ := classifyTestManager(t)
+	v := &VM{ID: "i-stopped", Status: StateStopped, InstanceType: "t3.micro"}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	assert.Empty(t, toLaunch, "stopped must not be queued for relaunch")
+	assert.NotNil(t, store.stopped[v.ID], "stopped must be migrated to the shared bucket")
+	_, stillLocal := m.Get(v.ID)
+	assert.False(t, stillLocal, "stopped must be removed from the local map after migration")
+}
+
+func TestClassifyRestoredInstances_UnknownTypeMarkedUnschedulable(t *testing.T) {
+	m, store, _ := classifyTestManager(t)
+	v := &VM{
+		ID:           "i-unknown",
+		Status:       StateRunning,
+		InstanceType: "z99.nonexistent",
+		Instance:     &ec2.Instance{},
+	}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	assert.Empty(t, toLaunch, "unschedulable type must not relaunch on this node")
+	assert.Equal(t, StateStopped, v.Status)
+	assert.NotNil(t, store.stopped[v.ID])
+	require.NotNil(t, v.Instance.StateReason)
+	assert.Equal(t, "Server.InsufficientInstanceCapacity", *v.Instance.StateReason.Code)
+}
+
+func TestClassifyRestoredInstances_StoppingFinalizesToStopped(t *testing.T) {
+	m, store, _ := classifyTestManager(t)
+	v := &VM{ID: "i-stopping", Status: StateStopping, InstanceType: "t3.micro"}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	assert.Empty(t, toLaunch)
+	assert.Equal(t, StateStopped, v.Status, "Stopping → Stopped finalisation when QEMU is gone")
+	assert.NotNil(t, store.stopped[v.ID])
+}
+
+func TestClassifyRestoredInstances_ShuttingDownFinalizesToTerminated(t *testing.T) {
+	m, store, _ := classifyTestManager(t)
+	v := &VM{ID: "i-shutdown", Status: StateShuttingDown, InstanceType: "t3.micro"}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	assert.Empty(t, toLaunch)
+	assert.Equal(t, StateTerminated, v.Status, "ShuttingDown → Terminated finalisation when QEMU is gone")
+	assert.NotNil(t, store.terminated[v.ID])
+}
+
+// TestClassifyRestoredInstances_RunningWithDeadPidQueuedForRelaunch
+// covers the no-PID-file path: Status was Running at shutdown but no
+// QEMU process is alive. classify must reset Status to Pending and
+// queue the instance for relaunch.
+//
+// Live-PID coverage (reconnect success/failure) is intentionally not
+// driven from a vm-package unit test — the SIGKILL path inside
+// killOrphanedQEMU and the long blocking PID-file polls in
+// MarkFailed's goroutine make process-spawn fakes too risky for a
+// hermetic test. The reconnect-failure regression is covered
+// separately by a daemon-side test against a real QEMU.
+func TestClassifyRestoredInstances_RunningWithDeadPidQueuedForRelaunch(t *testing.T) {
+	m, _, rc := classifyTestManager(t)
+	v := &VM{
+		ID:           "i-running-dead",
+		Status:       StateRunning,
+		InstanceType: "t3.micro",
+		Instance:     &ec2.Instance{},
+	}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	require.Len(t, toLaunch, 1)
+	assert.Same(t, v, toLaunch[0])
+	assert.Equal(t, StatePending, v.Status, "QEMU exited → reset to Pending so the relaunch path accepts it")
+	require.NotNil(t, v.Instance.LaunchTime, "LaunchTime must be reset for the pending watchdog")
+	assert.Equal(t, 1, rc.allocateCount("t3.micro"),
+		"resources must be re-allocated before queuing for relaunch")
+}
+
+func TestClassifyRestoredInstances_PendingQueuedForRelaunch(t *testing.T) {
+	m, _, rc := classifyTestManager(t)
+	v := &VM{
+		ID:           "i-pending",
+		Status:       StatePending,
+		InstanceType: "t3.micro",
+		Instance:     &ec2.Instance{},
+	}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	require.Len(t, toLaunch, 1)
+	assert.Equal(t, StatePending, v.Status, "Pending must remain Pending — no transitional reset")
+	require.NotNil(t, v.Instance.LaunchTime, "LaunchTime must be reset for the pending watchdog")
+	assert.Equal(t, 1, rc.allocateCount("t3.micro"),
+		"resources must be re-allocated before queuing for relaunch")
+}
+
+func TestClassifyRestoredInstances_AllocateFailureMarksUnschedulable(t *testing.T) {
+	m, store, rc := classifyTestManager(t)
+	rc.allocateErr = errors.New("no capacity")
+
+	v := &VM{
+		ID:           "i-no-cap",
+		Status:       StateRunning,
+		InstanceType: "t3.micro",
+		Instance:     &ec2.Instance{},
+	}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	assert.Empty(t, toLaunch, "allocate failure must demote the instance, not relaunch it")
+	assert.Equal(t, StateStopped, v.Status)
+	assert.NotNil(t, store.stopped[v.ID])
+	require.NotNil(t, v.Instance.StateReason)
+	assert.Equal(t, "Server.InsufficientInstanceCapacity", *v.Instance.StateReason.Code)
+}
+
+// finalizeRevertStateStore fails every shared-bucket write and the
+// running-state save so finalizeTransitionalRestore exercises the
+// revert-on-write-failure path. Embeds fakeStateStore so we still
+// satisfy the full interface.
+type finalizeRevertStateStore struct {
+	*fakeStateStore
+}
+
+func (finalizeRevertStateStore) WriteStoppedInstance(string, *VM) error {
+	return errors.New("stopped write failed")
+}
+func (finalizeRevertStateStore) WriteTerminatedInstance(string, *VM) error {
+	return errors.New("terminated write failed")
+}
+func (finalizeRevertStateStore) SaveRunningState(string, map[string]*VM) error {
+	return errors.New("save failed")
+}
+
+// TestFinalizeTransitionalRestore_WriteFailureReverts covers the
+// last-resort revert in finalizeTransitionalRestore: when both the
+// shared-bucket migration and the running-state save fail, the
+// instance's Status must revert to its pre-classify value so the next
+// daemon restart retries the same path. Without the revert the
+// instance would be left in an unstable in-memory state that no
+// startup branch knows how to recover.
+func TestFinalizeTransitionalRestore_WriteFailureReverts(t *testing.T) {
+	m := NewManager()
+	m.SetDeps(Deps{
+		NodeID:     "test-node",
+		StateStore: finalizeRevertStateStore{newFakeStateStore()},
+	})
+
+	v := &VM{ID: "i-revert", Status: StateStopping, InstanceType: "t3.micro"}
+	m.Insert(v)
+
+	ok := m.finalizeTransitionalRestore(v)
+
+	assert.True(t, ok, "finalizeTransitionalRestore returns true on the revert path so the caller continues")
+	assert.Equal(t, StateStopping, v.Status,
+		"both the migrate and the save failed — Status must revert so the next restart retries cleanly")
+}
+
+// markerStateStore captures the LoadRunningState call for assertions
+// against Restore's clean-shutdown handling. Returns the canned snapshot.
+type markerStateStore struct {
+	*fakeStateStore
+
+	loadCalls atomic.Int32
+	snapshot  map[string]*VM
+}
+
+func (s *markerStateStore) LoadRunningState(string) (map[string]*VM, error) {
+	s.loadCalls.Add(1)
+	out := make(map[string]*VM, len(s.snapshot))
+	maps.Copy(out, s.snapshot)
+	return out, nil
+}
+
+// TestRestore_HappyPath exercises Restore's clean-shutdown branch end-
+// to-end with a snapshot that classifies cleanly without entering
+// relaunch fan-out. Asserts: the marker callback is consulted, the
+// state store is loaded, terminated/stopped instances migrate to
+// their shared buckets, and the running-state snapshot is persisted.
+func TestRestore_HappyPath(t *testing.T) {
+	t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
+
+	store := &markerStateStore{
+		fakeStateStore: newFakeStateStore(),
+		snapshot: map[string]*VM{
+			"i-term":    {ID: "i-term", Status: StateTerminated, InstanceType: "t3.micro"},
+			"i-stopped": {ID: "i-stopped", Status: StateStopped, InstanceType: "t3.micro"},
+		},
+	}
+
+	markerConsumed := atomic.Bool{}
+	m := NewManager()
+	m.SetDeps(Deps{
+		NodeID:     "test-node",
+		StateStore: store,
+		ConsumeCleanShutdownMarker: func() bool {
+			markerConsumed.Store(true)
+			return true
+		},
+		InstanceTypes: fakeInstanceTypeResolver{"t3.micro": {VCPUs: 1, MemoryMiB: 1024, Architecture: "x86_64"}},
+		Resources:     newFakeResourceController(),
+	})
+
+	m.Restore()
+
+	assert.True(t, markerConsumed.Load(), "Restore must consult ConsumeCleanShutdownMarker")
+	assert.Equal(t, int32(1), store.loadCalls.Load(), "Restore must call LoadRunningState exactly once")
+	assert.NotNil(t, store.terminated["i-term"])
+	assert.NotNil(t, store.stopped["i-stopped"])
+	saved, ok := store.saved["test-node"]
+	require.True(t, ok, "Restore must persist the running snapshot at the end")
+	assert.Empty(t, saved, "running snapshot must be empty after both instances migrated to shared buckets")
+}
+
+func TestRestore_NoStateStore_NoCrash(t *testing.T) {
+	t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
+	m := NewManager()
+	m.SetDeps(Deps{
+		NodeID:                     "test-node",
+		ConsumeCleanShutdownMarker: func() bool { return true },
+	})
+
+	// loadRunningState surfaces "StateStore not wired"; Restore must
+	// log and return without panic.
+	require.NotPanics(t, m.Restore)
+}

spinifex/vm/shutdown_test.go

@@ -434,3 +434,308 @@ func TestCleanup_NoGPU_StillInvokesReleaseGPU(t *testing.T) {
 		t.Fatalf("ReleaseGPU calls across stop+terminate: got %d, want 2", got)
 	}
 }
+
+// TestTransitionWithPrecheck_Raced_WrapsErrInvalidTransition covers the
+// raced branch in transitionWithPrecheck: the static precheck succeeded
+// (transition was valid at that moment) but Deps.TransitionState
+// returned an error AND the in-memory status is no longer at target,
+// meaning a concurrent goroutine flipped the instance away. The
+// returned error must wrap ErrInvalidTransition so the daemon's
+// handler can map it to AWS IncorrectInstanceState (and so
+// daemon_handlers_instance.go's slog level keying on errors.Is
+// classifies the error correctly).
+func TestTransitionWithPrecheck_Raced_WrapsErrInvalidTransition(t *testing.T) {
+	persistErr := errors.New("kv put failed")
+	m := NewManagerWithDeps(Deps{
+		TransitionState: func(_ *VM, _ InstanceState) error {
+			// Return error without mutating status so the in-memory
+			// state stays at the pre-transition value, matching the
+			// "concurrent terminate beat us to it" scenario.
+			return persistErr
+		},
+	})
+
+	instance := &VM{ID: "i-raced", Status: StatePending}
+	m.Insert(instance)
+
+	err := m.transitionWithPrecheck(instance, StateRunning)
+
+	require.Error(t, err)
+	require.ErrorIs(t, err, ErrInvalidTransition,
+		"raced post-precheck failure must wrap ErrInvalidTransition so callers map to IncorrectInstanceState")
+	assert.NotErrorIs(t, err, persistErr,
+		"the original persistence error is intentionally hidden — caller surface is ErrInvalidTransition")
+	assert.Equal(t, StatePending, m.Status(instance),
+		"raced branch implies status was not flipped to target")
+}
+
+// TestTransitionWithPrecheck_PersistenceFailure_PassesThroughError is
+// the contrast: TransitionState returned an error but the in-memory
+// status DID reach target (the memory mutation succeeded; only the
+// persistence side failed). The original error must propagate as-is
+// without an ErrInvalidTransition wrap, so callers can distinguish a
+// transient persistence retry from a real transition rejection.
+func TestTransitionWithPrecheck_PersistenceFailure_PassesThroughError(t *testing.T) {
+	persistErr := errors.New("kv put failed")
+	var m *Manager
+	m = NewManagerWithDeps(Deps{
+		TransitionState: func(v *VM, target InstanceState) error {
+			// Memory mutation succeeded; only persistence failed.
+			m.Inspect(v, func(vv *VM) { vv.Status = target })
+			return persistErr
+		},
+	})
+
+	instance := &VM{ID: "i-persist-fail", Status: StatePending}
+	m.Insert(instance)
+
+	err := m.transitionWithPrecheck(instance, StateRunning)
+
+	require.Error(t, err)
+	require.ErrorIs(t, err, persistErr,
+		"persistence-only failure must surface the original error so callers can retry/log it specifically")
+	assert.NotErrorIs(t, err, ErrInvalidTransition,
+		"persistence failure on a valid transition is not a transition error")
+	assert.Equal(t, StateRunning, m.Status(instance),
+		"persistence-failure branch implies status reached target before error returned")
+}
+
+// TestTransitionWithPrecheck_InvalidInitialTransition_WrapsErrInvalidTransition
+// covers the static precheck rejecting an illegal transition before
+// TransitionState is ever called. The error must wrap
+// ErrInvalidTransition. Pairs with the raced branch — both surfaces
+// look the same to the caller, which is the design.
+func TestTransitionWithPrecheck_InvalidInitialTransition_WrapsErrInvalidTransition(t *testing.T) {
+	called := false
+	m := NewManagerWithDeps(Deps{
+		TransitionState: func(*VM, InstanceState) error {
+			called = true
+			return nil
+		},
+	})
+
+	instance := &VM{ID: "i-bad-trans", Status: StateTerminated}
+	m.Insert(instance)
+
+	err := m.transitionWithPrecheck(instance, StateRunning)
+
+	require.ErrorIs(t, err, ErrInvalidTransition)
+	assert.False(t, called, "Deps.TransitionState must not run when precheck rejects")
+	assert.Equal(t, StateTerminated, m.Status(instance), "rejected precheck must not mutate status")
+}
+
+// failingTerminatedStateStore overrides only WriteTerminatedInstance so
+// finalizeTerminated's persistence-failure path can be exercised.
+type failingTerminatedStateStore struct {
+	*fakeStateStore
+
+	err error
+}
+
+func (f *failingTerminatedStateStore) WriteTerminatedInstance(string, *VM) error {
+	return f.err
+}
+
+// reclaimingTerminatedStore reclaims the slot under the same id during
+// the terminated-bucket write so the subsequent DeleteIf in
+// finalizeTerminated finds a different pointer and returns false.
+type reclaimingTerminatedStore struct {
+	*fakeStateStore
+
+	m *Manager
+}
+
+func (r *reclaimingTerminatedStore) WriteTerminatedInstance(id string, v *VM) error {
+	if err := r.fakeStateStore.WriteTerminatedInstance(id, v); err != nil {
+		return err
+	}
+	r.m.Insert(&VM{ID: id, Status: StateRunning})
+	return nil
+}
+
+// failingSaveRunningStore overrides only SaveRunningState so the
+// post-DeleteIf rollback path in finalizeTerminated can be exercised
+// without disrupting WriteTerminatedInstance.
+type failingSaveRunningStore struct {
+	*fakeStateStore
+
+	err error
+}
+
+func (f *failingSaveRunningStore) SaveRunningState(string, map[string]*VM) error {
+	return f.err
+}
+
+// terminateTestManager builds a Manager wired for Terminate end-to-end
+// tests with hooks counter, state store, transition recorder, and a
+// recording cleaner. Returns everything callers may need to assert on.
+func terminateTestManager(t *testing.T, store StateStore) (m *Manager, cleaner *recordingInstanceCleaner, rt *recordedTransitions, downCount *atomic.Int64, downIDs *[]string) {
+	t.Helper()
+	cleaner = &recordingInstanceCleaner{}
+	rt = &recordedTransitions{}
+	downCount = &atomic.Int64{}
+	downIDs = &[]string{}
+	var mu sync.Mutex
+	m = NewManager()
+	rt.bind(m)
+	m.SetDeps(Deps{
+		NodeID:          "test-node",
+		StateStore:      store,
+		VolumeMounter:   &fakeVolumeMounter{},
+		InstanceCleaner: cleaner,
+		TransitionState: rt.apply,
+		ShutdownSignal:  func() bool { return false },
+		Hooks: ManagerHooks{
+			OnInstanceDown: func(id string) {
+				downCount.Add(1)
+				mu.Lock()
+				*downIDs = append(*downIDs, id)
+				mu.Unlock()
+			},
+		},
+	})
+	return m, cleaner, rt, downCount, downIDs
+}
+
+// TestTerminate_NotFound covers the validation guard at the top of
+// Terminate. An unknown id must return ErrInstanceNotFound without
+// touching cleaners, transitions, or hooks.
+func TestTerminate_NotFound(t *testing.T) {
+	m, cleaner, rt, down, _ := terminateTestManager(t, newFakeStateStore())
+
+	err := m.Terminate("i-missing")
+
+	require.ErrorIs(t, err, ErrInstanceNotFound)
+	assert.Empty(t, cleaner.deleteVolumes)
+	assert.Empty(t, rt.snapshot())
+	assert.Zero(t, down.Load())
+}
+
+// TestTerminate_AlreadyShuttingDown_Idempotent covers the early-return
+// at shutdown.go:113-116. A concurrent failed-launch goroutine has
+// already transitioned the instance to StateShuttingDown and owns
+// cleanup; Terminate must return nil and do nothing — no second
+// terminateCleanup, no second transition, no second OnInstanceDown.
+func TestTerminate_AlreadyShuttingDown_Idempotent(t *testing.T) {
+	m, cleaner, rt, down, _ := terminateTestManager(t, newFakeStateStore())
+	v := &VM{ID: "i-already-shutting", Status: StateShuttingDown, Instance: &ec2.Instance{}}
+	m.Insert(v)
+
+	err := m.Terminate(v.ID)
+
+	require.NoError(t, err)
+	assert.Empty(t, cleaner.deleteVolumes, "terminateCleanup must not run on idempotent path")
+	assert.Empty(t, rt.snapshot(), "no transitions on idempotent path")
+	assert.Zero(t, down.Load(), "OnInstanceDown must not fire on idempotent path")
+	assert.Equal(t, StateShuttingDown, m.Status(v))
+}
+
+// TestTerminate_InvalidTransition covers transitionWithPrecheck rejecting
+// an illegal initial transition. StateTerminated → StateShuttingDown is
+// not in ValidTransitions, so the daemon-side handler can map this to
+// AWS IncorrectInstanceState.
+func TestTerminate_InvalidTransition(t *testing.T) {
+	m, cleaner, _, down, _ := terminateTestManager(t, newFakeStateStore())
+	v := &VM{ID: "i-already-terminated", Status: StateTerminated, Instance: &ec2.Instance{}}
+	m.Insert(v)
+
+	err := m.Terminate(v.ID)
+
+	require.ErrorIs(t, err, ErrInvalidTransition)
+	assert.Empty(t, cleaner.deleteVolumes, "terminateCleanup must not run when precheck rejects")
+	assert.Zero(t, down.Load())
+}
+
+// TestTerminate_Success_FiresHooksAndCleanupOnce is the happy-path
+// contract: a Running instance terminates through ShuttingDown →
+// Terminated, the full terminate-only cleanup chain runs once, the
+// terminated KV bucket is written, the instance leaves the local map,
+// OnInstanceDown fires exactly once, and the running-state snapshot
+// is persisted.
+func TestTerminate_Success_FiresHooksAndCleanupOnce(t *testing.T) {
+	store := newFakeStateStore()
+	m, cleaner, rt, down, downIDs := terminateTestManager(t, store)
+
+	v := &VM{ID: "i-terminate-ok", Status: StateRunning, InstanceType: "t3.micro", Instance: &ec2.Instance{}}
+	m.Insert(v)
+
+	require.NoError(t, m.Terminate(v.ID))
+
+	assert.Equal(t, []InstanceState{StateShuttingDown, StateTerminated}, rt.targets(v.ID))
+	assert.Equal(t, StateTerminated, m.Status(v))
+	assert.Equal(t, []string{v.ID}, cleaner.deleteVolumes, "terminate-only cleaner must run once")
+	assert.Equal(t, []string{v.ID}, cleaner.releasePublicIP)
+	assert.Equal(t, []string{v.ID}, cleaner.detachAndDeleteENI)
+	assert.Equal(t, []string{v.ID}, cleaner.removeFromPlacement)
+	assert.Equal(t, []string{v.ID}, cleaner.releaseGPU)
+	assert.Equal(t, int64(1), down.Load(), "OnInstanceDown must fire exactly once on success")
+	assert.Equal(t, []string{v.ID}, *downIDs)
+	require.NotNil(t, store.terminated[v.ID], "WriteTerminatedInstance must persist before DeleteIf")
+	_, stillInMap := m.Get(v.ID)
+	assert.False(t, stillInMap, "instance must leave the local map on successful Terminate")
+}
+
+// TestTerminate_WriteTerminatedFailure_PropagatesAndKeepsLocal covers
+// finalizeTerminated's persistence guard at shutdown.go:189-194. When
+// the terminated-bucket write fails the error must propagate and the
+// instance must stay in the local map so the next daemon restart
+// can retry — DeleteIf must not run, OnInstanceDown must not fire.
+func TestTerminate_WriteTerminatedFailure_PropagatesAndKeepsLocal(t *testing.T) {
+	persistErr := errors.New("kv put failed")
+	store := &failingTerminatedStateStore{fakeStateStore: newFakeStateStore(), err: persistErr}
+	m, _, _, down, _ := terminateTestManager(t, store)
+
+	v := &VM{ID: "i-write-fail", Status: StateRunning, InstanceType: "t3.micro", Instance: &ec2.Instance{}}
+	m.Insert(v)
+
+	err := m.Terminate(v.ID)
+
+	require.ErrorIs(t, err, persistErr)
+	assert.Zero(t, down.Load(), "OnInstanceDown must not fire when the terminated-bucket write fails")
+	_, stillInMap := m.Get(v.ID)
+	assert.True(t, stillInMap, "instance must remain in local map for retry on the next daemon restart")
+}
+
+// TestTerminate_DeleteIfReclaimed_NoHook covers the slot-reclaim branch
+// in finalizeTerminated (shutdown.go:197-201): a concurrent handler
+// inserts a fresh VM under the same id between WriteTerminatedInstance
+// and DeleteIf. Terminate must return nil without firing
+// OnInstanceDown — firing it would tear down the new instance's NATS
+// subscriptions.
+func TestTerminate_DeleteIfReclaimed_NoHook(t *testing.T) {
+	store := &reclaimingTerminatedStore{fakeStateStore: newFakeStateStore()}
+	m, _, _, down, _ := terminateTestManager(t, store)
+	store.m = m
+
+	v := &VM{ID: "i-reclaim-term", Status: StateRunning, InstanceType: "t3.micro", Instance: &ec2.Instance{}}
+	m.Insert(v)
+
+	require.NoError(t, m.Terminate(v.ID))
+
+	assert.Zero(t, down.Load(), "slot-reclaim must not fire OnInstanceDown")
+	got, ok := m.Get(v.ID)
+	require.True(t, ok, "the reclaiming inserter's pointer must still be in the map")
+	assert.NotSame(t, v, got, "DeleteIf must not have removed the reclaimer's VM")
+}
+
+// TestTerminate_WriteRunningStateFailure_RollbackInsert covers the
+// rollback path at shutdown.go:207-212. SaveRunningState fails after
+// DeleteIf has already removed the instance; finalizeTerminated must
+// re-insert it so the local map and the persisted snapshot stay
+// consistent. Terminate still returns nil because the state write is
+// best-effort at this layer.
+func TestTerminate_WriteRunningStateFailure_RollbackInsert(t *testing.T) {
+	saveErr := errors.New("save failed")
+	store := &failingSaveRunningStore{fakeStateStore: newFakeStateStore(), err: saveErr}
+	m, _, _, down, _ := terminateTestManager(t, store)
+
+	v := &VM{ID: "i-saverun-fail", Status: StateRunning, InstanceType: "t3.micro", Instance: &ec2.Instance{}}
+	m.Insert(v)
+
+	require.NoError(t, m.Terminate(v.ID))
+
+	assert.Equal(t, int64(1), down.Load(), "OnInstanceDown fires before the writeRunningState rollback")
+	got, ok := m.Get(v.ID)
+	require.True(t, ok, "writeRunningState failure must re-insert the instance into the local map")
+	assert.Same(t, v, got)
+}