From cde092cbfbb66b2a4c4cafa451f27431366163ec Mon Sep 17 00:00:00 2001 From: Julian Sommer <115516918+juliansommer@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:19:50 +1000 Subject: [PATCH] fix: speedup daemon tests --- spinifex/daemon/daemon_system_instance.go | 5 ++- .../daemon/daemon_system_instance_test.go | 35 ++++++++++++++++--- spinifex/daemon/disconnect_test.go | 9 +++-- .../daemon/instance_records_rekey_test.go | 4 ++- spinifex/utils/nats.go | 14 +++++++- 5 files changed, 57 insertions(+), 10 deletions(-) diff --git a/spinifex/daemon/daemon_system_instance.go b/spinifex/daemon/daemon_system_instance.go index 641727169..0017305b0 100644 --- a/spinifex/daemon/daemon_system_instance.go +++ b/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) } diff --git a/spinifex/daemon/daemon_system_instance_test.go b/spinifex/daemon/daemon_system_instance_test.go index f65b54a0c..d0b649e76 100644 --- a/spinifex/daemon/daemon_system_instance_test.go +++ b/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") + } } // --------------------------------------------------------------------------- diff --git a/spinifex/daemon/disconnect_test.go b/spinifex/daemon/disconnect_test.go index 8ecc61121..f9b48051c 100644 --- a/spinifex/daemon/disconnect_test.go +++ b/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()) diff --git a/spinifex/daemon/instance_records_rekey_test.go b/spinifex/daemon/instance_records_rekey_test.go index 4b1f3a8d7..fa9a3f1f7 100644 --- a/spinifex/daemon/instance_records_rekey_test.go +++ b/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): } } diff --git a/spinifex/utils/nats.go b/spinifex/utils/nats.go index 6d6f5a151..788032c79 100644 --- a/spinifex/utils/nats.go +++ b/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 {