ASCII-808: Linter ensures fxutil.Run is tested using TestRun

cmd/otel-agent/main_test.go

@@ -0,0 +1,20 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+//go:build otlp
+
+package main
+
+import (
+	"testing"
+
+	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
+)
+
+func TestFxRun(t *testing.T) {
+	// TODO: (components) "missing type: *aggregator.AgentDemultiplexer"
+	t.SkipNow()
+	fxutil.TestOneShot(t, main)
+}

cmd/security-agent/main_windows_test.go

@@ -0,0 +1,25 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+//go:build windows
+
+package main
+
+import (
+	"context"
+	"testing"
+
+	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
+	"github.com/stretchr/testify/require"
+)
+
+func TestFxRun(t *testing.T) {
+	svc := service{}
+	ctx := context.Background()
+	fxutil.TestOneShot(t, func() {
+		err := svc.Run(ctx)
+		require.NoError(t, err)
+	})
+}

cmd/serverless-init/main_test.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/DataDog/datadog-agent/pkg/config"
 	"github.com/DataDog/datadog-agent/pkg/serverless/logs"
+	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
 )
 
 func setupTest() {
@@ -43,3 +44,7 @@ func TestTagsSetup(t *testing.T) {
 	assert.Subset(t, metricAgent.GetExtraTags(), allTags)
 	assert.Subset(t, logs.GetLogsTags(), allTags)
 }
+
+func TestFxApp(t *testing.T) {
+	fxutil.TestOneShot(t, main)
+}

cmd/serverless/main_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/DataDog/datadog-agent/pkg/serverless/daemon"
+	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
 )
 
 func TestDaemonStopOnTerminationSignals(t *testing.T) {
@@ -52,3 +53,7 @@ func TestDaemonStopOnTerminationSignals(t *testing.T) {
 		})
 	}
 }
+
+func TestFxApp(t *testing.T) {
+	fxutil.TestOneShot(t, main)
+}

cmd/system-probe/subcommands/run/command.go

@@ -11,6 +11,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+
 	//nolint:revive // TODO(EBPF) Fix revive linter
 	_ "net/http/pprof"
 	"os"
@@ -161,48 +162,10 @@ func run(log log.Component, _ config.Component, statsd compstatsd.Component, tel
 func StartSystemProbeWithDefaults(ctxChan <-chan context.Context) (<-chan error, error) {
 	errChan := make(chan error)
 
-	// run startSystemProbe in an app, so that the log and config components get initialized
+	// run startSystemProbe in the background
 	go func() {
-		err := fxutil.OneShot(
-			func(log log.Component, config config.Component, statsd compstatsd.Component, telemetry telemetry.Component, sysprobeconfig sysprobeconfig.Component, rcclient rcclient.Component) error {
-				defer StopSystemProbeWithDefaults()
-				err := startSystemProbe(&cliParams{GlobalParams: &command.GlobalParams{}}, log, statsd, telemetry, sysprobeconfig, rcclient)
-				if err != nil {
-					return err
-				}
-
-				// notify outer that startAgent finished
-				errChan <- err
-				// wait for context
-				ctx := <-ctxChan
-
-				// Wait for stop signal
-				select {
-				case <-signals.Stopper:
-					log.Info("Received stop command, shutting down...")
-				case <-signals.ErrorStopper:
-					_ = log.Critical("The Agent has encountered an error, shutting down...")
-				case <-ctx.Done():
-					log.Info("Received stop from service manager, shutting down...")
-				}
-
-				return nil
-			},
-			// no config file path specification in this situation
-			fx.Supply(config.NewAgentParams("", config.WithConfigMissingOK(true))),
-			fx.Supply(sysprobeconfigimpl.NewParams(sysprobeconfigimpl.WithSysProbeConfFilePath(""))),
-			fx.Supply(logimpl.ForDaemon("SYS-PROBE", "log_file", common.DefaultLogFile)),
-			rcclient.Module(),
-			config.Module(),
-			telemetry.Module(),
-			compstatsd.Module(),
-			sysprobeconfigimpl.Module(),
-			// use system-probe config instead of agent config for logging
-			fx.Provide(func(lc fx.Lifecycle, params logimpl.Params, sysprobeconfig sysprobeconfig.Component) (log.Component, error) {
-				return logimpl.NewLogger(lc, params, sysprobeconfig)
-			}),
-		)
-		// notify caller that fx.OneShot is done
+		err := runSystemProbe(ctxChan, errChan)
+		// notify main routine that this is done, so cleanup can happen
 		errChan <- err
 	}()
 
@@ -217,6 +180,48 @@ func StartSystemProbeWithDefaults(ctxChan <-chan context.Context) (<-chan error,
 	return errChan, nil
 }
 
+func runSystemProbe(ctxChan <-chan context.Context, errChan chan error) error {
+	return fxutil.OneShot(
+		func(log log.Component, config config.Component, statsd compstatsd.Component, telemetry telemetry.Component, sysprobeconfig sysprobeconfig.Component, rcclient rcclient.Component) error {
+			defer StopSystemProbeWithDefaults()
+			err := startSystemProbe(&cliParams{GlobalParams: &command.GlobalParams{}}, log, statsd, telemetry, sysprobeconfig, rcclient)
+			if err != nil {
+				return err
+			}
+
+			// notify outer that startAgent finished
+			errChan <- err
+			// wait for context
+			ctx := <-ctxChan
+
+			// Wait for stop signal
+			select {
+			case <-signals.Stopper:
+				log.Info("Received stop command, shutting down...")
+			case <-signals.ErrorStopper:
+				_ = log.Critical("The Agent has encountered an error, shutting down...")
+			case <-ctx.Done():
+				log.Info("Received stop from service manager, shutting down...")
+			}
+
+			return nil
+		},
+		// no config file path specification in this situation
+		fx.Supply(config.NewAgentParams("", config.WithConfigMissingOK(true))),
+		fx.Supply(sysprobeconfigimpl.NewParams(sysprobeconfigimpl.WithSysProbeConfFilePath(""))),
+		fx.Supply(logimpl.ForDaemon("SYS-PROBE", "log_file", common.DefaultLogFile)),
+		rcclient.Module(),
+		config.Module(),
+		telemetry.Module(),
+		compstatsd.Module(),
+		sysprobeconfigimpl.Module(),
+		// use system-probe config instead of agent config for logging
+		fx.Provide(func(lc fx.Lifecycle, params logimpl.Params, sysprobeconfig sysprobeconfig.Component) (log.Component, error) {
+			return logimpl.NewLogger(lc, params, sysprobeconfig)
+		}),
+	)
+}
+
 // StopSystemProbeWithDefaults is a temporary way for other packages to use stopAgent.
 func StopSystemProbeWithDefaults() {
 	stopSystemProbe(&cliParams{GlobalParams: &command.GlobalParams{}})

cmd/system-probe/subcommands/run/command_test.go

@@ -6,6 +6,7 @@
 package run
 
 import (
+	"context"
 	_ "net/http/pprof"
 	"testing"
 
@@ -20,3 +21,11 @@ func TestRunCommand(t *testing.T) {
 		run,
 		func() {})
 }
+
+func TestStartSystemProbe(t *testing.T) {
+	fxutil.TestOneShot(t, func() {
+		ctxChan := make(<-chan context.Context)
+		errChan := make(chan error)
+		runSystemProbe(ctxChan, errChan)
+	})
+}

cmd/systray/command/command_test.go

@@ -0,0 +1,20 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+//go:build windows
+
+package command
+
+import (
+	"testing"
+
+	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
+)
+
+func TestFxRunCommand(t *testing.T) {
+	cmd := MakeCommand()
+	fxutil.TestRun(t, func() error {
+		return cmd.Execute()
+	})
+}

cmd/trace-agent/subcommands/run/command.go

@@ -39,7 +39,7 @@ func MakeCommand(globalParamsGetter func() *subcommands.GlobalParams) *cobra.Com
 		Long:  `The Datadog trace-agent aggregates, samples, and forwards traces to datadog submitted by tracers loaded into your application.`,
 		RunE: func(*cobra.Command, []string) error {
 			cliParams.GlobalParams = globalParamsGetter()
-			return runTraceAgent(cliParams, cliParams.ConfPath)
+			return runTraceAgentCommand(cliParams, cliParams.ConfPath)
 		},
 	}
 
@@ -58,7 +58,7 @@ func setParamFlags(cmd *cobra.Command, cliParams *RunParams) {
 	setOSSpecificParamFlags(cmd, cliParams)
 }
 
-func runFx(ctx context.Context, cliParams *RunParams, defaultConfPath string) error {
+func runTraceAgentProcess(ctx context.Context, cliParams *RunParams, defaultConfPath string) error {
 	if cliParams.ConfPath == "" {
 		cliParams.ConfPath = defaultConfPath
 	}

cmd/trace-agent/subcommands/run/command_nix.go

@@ -30,6 +30,6 @@ type RunParams struct {
 //nolint:revive // TODO(APM) Fix revive linter
 func setOSSpecificParamFlags(cmd *cobra.Command, cliParams *RunParams) {}
 
-func runTraceAgent(cliParams *RunParams, defaultConfPath string) error {
-	return runFx(context.Background(), cliParams, defaultConfPath)
+func runTraceAgentCommand(cliParams *RunParams, defaultConfPath string) error {
+	return runTraceAgentProcess(context.Background(), cliParams, defaultConfPath)
 }

cmd/trace-agent/subcommands/run/command_test.go

@@ -0,0 +1,23 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+package run
+
+import (
+	"context"
+	"testing"
+
+	"github.com/DataDog/datadog-agent/cmd/trace-agent/subcommands"
+	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
+)
+
+func TestFxRun(t *testing.T) {
+	fxutil.TestRun(t, func() error {
+		ctx := context.Background()
+		cliParams := RunParams{GlobalParams: &subcommands.GlobalParams{}}
+		defaultConfPath := ""
+		return runTraceAgentProcess(ctx, &cliParams, defaultConfPath)
+	})
+}

cmd/trace-agent/subcommands/run/command_windows.go

@@ -38,12 +38,12 @@ func setOSSpecificParamFlags(cmd *cobra.Command, cliParams *RunParams) {
 		"runs the trace-agent in debug mode.")
 }
 
-func runTraceAgent(cliParams *RunParams, defaultConfPath string) error {
+func runTraceAgentCommand(cliParams *RunParams, defaultConfPath string) error {
 	if !cliParams.Foreground {
 		if servicemain.RunningAsWindowsService() {
 			servicemain.Run(&service{cliParams: cliParams, defaultConfPath: defaultConfPath})
 			return nil
 		}
 	}
-	return runFx(context.Background(), cliParams, defaultConfPath)
+	return runTraceAgentProcess(context.Background(), cliParams, defaultConfPath)
 }

cmd/trace-agent/subcommands/run/service_windows.go

@@ -28,5 +28,5 @@ func (s *service) Init() error {
 }
 
 func (s *service) Run(ctx context.Context) error {
-	return runFx(ctx, s.cliParams, s.defaultConfPath)
+	return runTraceAgentProcess(ctx, s.cliParams, s.defaultConfPath)
 }

pkg/util/fxutil/oneshot.go

@@ -11,10 +11,6 @@ import (
 	"go.uber.org/fx"
 )
 
-// oneShotTestOverride allows TestOneShotSubcommand to override the OneShot
-// function.  It is always nil in production.
-var oneShotTestOverride func(interface{}, []fx.Option) error
-
 // OneShot runs the given function in an fx.App using the supplied options.
 // The function's arguments are supplied by Fx and can be any provided type.
 // The function must return `error` or nothing.
@@ -23,8 +19,8 @@ var oneShotTestOverride func(interface{}, []fx.Option) error
 // immediately shuts down.  This is typically used for command-line tools like
 // `agent status`.
 func OneShot(oneShotFunc interface{}, opts ...fx.Option) error {
-	if oneShotTestOverride != nil {
-		return oneShotTestOverride(oneShotFunc, opts)
+	if fxAppTestOverride != nil {
+		return fxAppTestOverride(oneShotFunc, opts)
 	}
 
 	// Use a delayed Fx invocation to capture arguments for oneShotFunc during

pkg/util/fxutil/run.go

@@ -16,6 +16,10 @@ import (
 // This differs from fx.App#Run in that it returns errors instead of exiting
 // the process.
 func Run(opts ...fx.Option) error {
+	if fxAppTestOverride != nil {
+		return fxAppTestOverride(func() {}, opts)
+	}
+
 	opts = append(opts, FxLoggingOption())
 	// Temporarily increase timeout for all fxutil.Run calls until we can better characterize our
 	// start time requirements. Prepend to opts so individual calls can override the timeout.