fix: wait for agent to be up in TestAgent_Lifecycle/ShutdownScriptOnce

spikecurtis · spikecurtis · commit 82bdf7f827a8 · 2025-04-14T16:10:47.000+04:00
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -1650,8 +1650,10 @@ func TestAgent_Lifecycle(t *testing.T) {
 	t.Run("ShutdownScriptOnce", func(t *testing.T) {
 		t.Parallel()
 		logger := testutil.Logger(t)
+		ctx := testutil.Context(t, testutil.WaitMedium)
 		expected := "this-is-shutdown"
 		derpMap, _ := tailnettest.RunDERPAndSTUN(t)
+		statsCh := make(chan *proto.Stats, 50)
 
 		client := agenttest.NewClient(t,
 			logger,
@@ -1670,7 +1672,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 					RunOnStop: true,
 				}},
 			},
-			make(chan *proto.Stats, 50),
+			statsCh,
 			tailnet.NewCoordinator(logger),
 		)
 		defer client.Close()
@@ -1695,6 +1697,11 @@ func TestAgent_Lifecycle(t *testing.T) {
 			return len(content) > 0 // something is in the startup log file
 		}, testutil.WaitShort, testutil.IntervalMedium)
 
+		// In order to avoid shutting down the agent before it is fully started and triggering
+		// errors, we'll wait until the agent is fully up. It's a bit hokey, but among the last things the agent starts
+		// is the stats reporting, so getting a stats report is a good indication the agent is fully up.
+		_ = testutil.RequireRecvCtx(ctx, t, statsCh)
+
 		err := agent.Close()
 		require.NoError(t, err, "agent should be closed successfully")
 
