fix: track goroutine for test + use separate context for reporting

DanielleMaywood · DanielleMaywood · commit 13889bf11531 · 2024-09-23T08:50:03.000Z
diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go
@@ -338,7 +338,7 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript,
 			return
 		}
 
-		go func() {
+		r.trackCommandGoroutine(func() {
 			var stage proto.Timing_Stage
 			switch option {
 			case ExecuteStartScripts:
@@ -349,7 +349,11 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript,
 				stage = proto.Timing_CRON
 			}
 
-			_, err = r.scriptCompleted(ctx, &proto.WorkspaceAgentScriptCompletedRequest{
+			reportTimeout := 30 * time.Second
+			reportCtx, cancel := context.WithTimeout(context.Background(), reportTimeout)
+			defer cancel()
+
+			_, err = r.scriptCompleted(reportCtx, &proto.WorkspaceAgentScriptCompletedRequest{
 				Timing: &proto.Timing{
 					ScriptId: script.ID[:],
 					Start:    timestamppb.New(start),
@@ -362,7 +366,7 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript,
 			if err != nil {
 				logger.Error(ctx, fmt.Sprintf("reporting script completed: %s", err.Error()))
 			}
-		}()
+		})
 	}()
 
 	err = cmd.Start()
diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go
@@ -120,7 +120,6 @@ func TestScriptReportsTiming(t *testing.T) {
 	runner := setup(t, func(uuid2 uuid.UUID) agentscripts.ScriptLogger {
 		return fLogger
 	})
-	defer runner.Close()
 
 	aAPI := agenttest.NewFakeAgentAPI(t, slogtest.Make(t, nil), nil, nil)
 	err := runner.Init([]codersdk.WorkspaceAgentScript{{
@@ -129,7 +128,8 @@ func TestScriptReportsTiming(t *testing.T) {
 		Script:      "echo hello",
 	}}, aAPI.ScriptCompleted)
 	require.NoError(t, err)
-	require.NoError(t, runner.Execute(context.Background(), agentscripts.ExecuteAllScripts))
+	require.NoError(t, runner.Execute(ctx, agentscripts.ExecuteAllScripts))
+	runner.Close()
 
 	log := testutil.RequireRecvCtx(ctx, t, fLogger.logs)
 	require.Equal(t, "hello", log.Output)
