fix: speedup daemon tests

spinifex/daemon/daemon_system_instance.go

@@ -731,6 +731,9 @@ func (d *Daemon) cleanupFailedSystemInstance(instance *vm.VM, _ *ec2.InstanceTyp
 	d.vmMgr.MarkFailed(context.Background(), instance, "system_instance_launch_failed")
 }
 
+// waitPollSleep is the poll delay seam used by WaitForSystemInstance; tests override it.
+var waitPollSleep = time.Sleep
+
 // WaitForSystemInstance polls until the instance reaches running state or times out.
 func (d *Daemon) WaitForSystemInstance(instanceID string, timeout time.Duration) error {
 	deadline := time.Now().Add(timeout)
@@ -749,7 +752,7 @@ func (d *Daemon) WaitForSystemInstance(instanceID string, timeout time.Duration)
 			return fmt.Errorf("instance %s in terminal state: %s", instanceID, status)
 		}
 
-		time.Sleep(500 * time.Millisecond)
+		waitPollSleep(500 * time.Millisecond)
 	}
 	return fmt.Errorf("instance %s did not reach running state within %s", instanceID, timeout)
 }

spinifex/daemon/daemon_system_instance_test.go

@@ -62,29 +62,54 @@ func TestWaitForSystemInstance_ErrorState(t *testing.T) {
 	}
 }
 
+// stubWaitPollSleep swaps the WaitForSystemInstance poll delay for fn so tests
+// drive the loop instead of waiting out 500ms per iteration. Restored via t.Cleanup.
+func stubWaitPollSleep(t *testing.T, fn func(time.Duration)) {
+	t.Helper()
+	prev := waitPollSleep
+	waitPollSleep = fn
+	t.Cleanup(func() { waitPollSleep = prev })
+}
+
 func TestWaitForSystemInstance_TransitionsToRunning(t *testing.T) {
 	inst := &vm.VM{ID: "i-pending", Status: vm.StateProvisioning}
 	d := newDaemonWithVMs(inst)
 
-	// Transition to running after a short delay
-	go func() {
-		time.Sleep(200 * time.Millisecond)
+	// Transition to running from inside the poll delay, so the second iteration
+	// observes it without the loop waiting on wall-clock time.
+	polls := 0
+	stubWaitPollSleep(t, func(time.Duration) {
+		polls++
 		d.vmMgr.UpdateState(inst.ID, func(v *vm.VM) { v.Status = vm.StateRunning })
-	}()
+	})
 
 	err := d.WaitForSystemInstance("i-pending", 2*time.Second)
 	if err != nil {
 		t.Fatalf("expected instance to reach running, got: %v", err)
 	}
+	if polls != 1 {
+		t.Fatalf("expected 1 poll delay before running was observed, got %d", polls)
+	}
 }
 
 func TestWaitForSystemInstance_Timeout(t *testing.T) {
 	d := newDaemonWithVMs(&vm.VM{ID: "i-stuck", Status: vm.StateProvisioning})
 
-	err := d.WaitForSystemInstance("i-stuck", 600*time.Millisecond)
+	// Sleep a token amount so the deadline is still reached by elapsed time
+	// rather than by spinning the loop.
+	polls := 0
+	stubWaitPollSleep(t, func(time.Duration) {
+		polls++
+		time.Sleep(time.Millisecond)
+	})
+
+	err := d.WaitForSystemInstance("i-stuck", 5*time.Millisecond)
 	if err == nil {
 		t.Fatal("expected timeout error")
 	}
+	if polls == 0 {
+		t.Fatal("expected the poll loop to run at least once before timing out")
+	}
 }
 
 // ---------------------------------------------------------------------------

spinifex/daemon/disconnect_test.go

@@ -19,6 +19,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// testReconnectWait shortens the client's per-server reconnect delay from the
+// 1s production default. Tests that kill and restart a server otherwise wait
+// out that delay before the reconnect handler fires.
+var testReconnectWait = utils.WithReconnectWait(50 * time.Millisecond)
+
 // TestDaemonModeDefaultStandalone — Mode() returns standalone before Start().
 func TestDaemonModeDefaultStandalone(t *testing.T) {
 	clusterCfg := &config.ClusterConfig{
@@ -45,7 +50,7 @@ func TestDaemonModeFlipsOnConnectAndDisconnect(t *testing.T) {
 	d, err := NewDaemon(clusterCfg)
 	require.NoError(t, err)
 
-	require.NoError(t, d.connectNATS())
+	require.NoError(t, d.connectNATS(testReconnectWait))
 	defer d.natsConn.Close()
 	assert.Equal(t, DaemonModeCluster, d.Mode())
 
@@ -66,7 +71,7 @@ func TestDaemonReconnectBumpsRetryCount(t *testing.T) {
 	d, err := NewDaemon(clusterCfg)
 	require.NoError(t, err)
 
-	require.NoError(t, d.connectNATS())
+	require.NoError(t, d.connectNATS(testReconnectWait))
 	defer d.natsConn.Close()
 	require.Equal(t, int64(0), d.NATSRetryCount())
 

spinifex/daemon/instance_records_rekey_test.go

@@ -150,9 +150,11 @@ func TestRecordKeys_SlashSeparatorWasNotWatchable(t *testing.T) {
 	_, err = kv.Put(context.Background(), slashPrefix+"i-1", []byte("{}"))
 	require.NoError(t, err)
 
+	// A watch that will deliver delivers in milliseconds — the server pushes on
+	// the Put. 200ms is a wide margin for the negative assertion.
 	select {
 	case entry := <-watcher.Updates():
 		t.Fatalf("the slash separator delivered a watch event for %q, so the re-key was not needed", entry.Key())
-	case <-time.After(2 * time.Second):
+	case <-time.After(200 * time.Millisecond):
 	}
 }

spinifex/utils/nats.go

@@ -60,8 +60,13 @@ func ConnectNATS(host, token, caCertPath string, opts ...RetryOption) (*nats.Con
 		o(&cfg)
 	}
 
+	reconnectWait := cfg.reconnectWait
+	if reconnectWait <= 0 {
+		reconnectWait = time.Second
+	}
+
 	natsOpts := []nats.Option{
-		nats.ReconnectWait(time.Second),
+		nats.ReconnectWait(reconnectWait),
 		nats.MaxReconnects(-1),
 		nats.DisconnectErrHandler(func(nc *nats.Conn, err error) {
 			slog.Warn("NATS disconnected", "err", err)
@@ -107,6 +112,7 @@ type retryConfig struct {
 	maxWait       time.Duration
 	retryDelay    time.Duration
 	maxRetryDelay time.Duration
+	reconnectWait time.Duration
 	onDisconnect  func(*nats.Conn, error)
 	onReconnect   func(*nats.Conn)
 	onAttemptErr  func(err error, attempt int)
@@ -132,6 +138,12 @@ func WithMaxRetryDelay(d time.Duration) RetryOption {
 	return func(c *retryConfig) { c.maxRetryDelay = d }
 }
 
+// WithReconnectWait overrides how long the client waits between reconnect
+// attempts to the same server (default 1s).
+func WithReconnectWait(d time.Duration) RetryOption {
+	return func(c *retryConfig) { c.reconnectWait = d }
+}
+
 // WithDisconnectHandler registers a callback invoked after the default disconnect log line.
 // Runs on a NATS goroutine; keep it non-blocking.
 func WithDisconnectHandler(fn func(*nats.Conn, error)) RetryOption {