patch: add lifecycle tests

 #133 must be merged first

- Currently if a panic occurs during our fx lifecycle, we log the panic and then ignore it.
- Instead, we want to return an error in order for fx to rollback and shutdown

cmd/xmidt-agent/main.go

@@ -29,6 +29,12 @@ const (
 	applicationName = "xmidt-agent"
 )
 
+var (
+	ErrLifecycleStartPanic    = errors.New("panic occured during fx's lifecycle Start")
+	ErrLifecycleStopPanic     = errors.New("panic occured during fx's lifecycle Stop")
+	ErrLifecycleShutdownPanic = errors.New("panic occured during fx's lifecycle Shutdown")
+)
+
 // These match what goreleaser provides.
 var (
 	commit  = "undefined"
@@ -60,6 +66,16 @@ type LifeCycleIn struct {
 // xmidtAgent is the main entry point for the program.  It is responsible for
 // setting up the dependency injection framework and returning the app object.
 func xmidtAgent(args []string) (*fx.App, error) {
+	app := fx.New(provideAppOptions(args))
+	if err := app.Err(); err != nil {
+		return nil, err
+	}
+
+	return app, nil
+}
+
+// provideAppOptions returns all fx options required to start the xmidt agent fx app.
+func provideAppOptions(args []string) fx.Option {
 	var (
 		gscfg *goschtalt.Config
 
@@ -70,7 +86,7 @@ func xmidtAgent(args []string) (*fx.App, error) {
 		cli *CLI
 	)
 
-	app := fx.New(
+	opts := fx.Options(
 		fx.Supply(cliArgs(args)),
 		fx.Populate(&g),
 		fx.Populate(&gscfg),
@@ -119,11 +135,7 @@ func xmidtAgent(args []string) (*fx.App, error) {
 		_ = os.WriteFile(cli.Graph, []byte(g), 0600)
 	}
 
-	if err := app.Err(); err != nil {
-		return nil, err
-	}
-
-	return app, nil
+	return opts
 }
 
 func main() {
@@ -222,12 +234,15 @@ func onStart(cred *credentials.Credentials, ws *websocket.Websocket, qos *qos.Ha
 		// err is set during a panic recovery in order to allow fx to rolling back
 		defer func() {
 			if r := recover(); nil != r {
-				logger.Error("stacktrace from panic", zap.String("stacktrace", string(debug.Stack())), zap.Any("panic", r))
-				err = errors.New("panic occured during fx lifecycle Start")
+				err = ErrLifecycleStartPanic
+				logger.Error("stacktrace from panic", zap.String("stacktrace", string(debug.Stack())), zap.Any("panic", r), zap.Error(err))
 			}
-
 		}()
 
+		if err = ctx.Err(); err != nil {
+			return err
+		}
+
 		if ws == nil {
 			logger.Debug("websocket disabled")
 			return err
@@ -251,13 +266,16 @@ func onStart(cred *credentials.Credentials, ws *websocket.Websocket, qos *qos.Ha
 func onStop(ws *websocket.Websocket, qos *qos.Handler, shutdowner fx.Shutdowner, cancels []func(), logger *zap.Logger) func(context.Context) error {
 	logger = logger.Named("on_stop")
 
-	return func(_ context.Context) error {
+	return func(context.Context) (err error) {
 		defer func() {
 			if r := recover(); nil != r {
-				logger.Error("stacktrace from panic", zap.String("stacktrace", string(debug.Stack())), zap.Any("panic", r))
+				err = ErrLifecycleStopPanic
+				fmt.Println(string(debug.Stack()))
+				logger.Error("stacktrace from panic", zap.String("stacktrace", string(debug.Stack())), zap.Any("panic", r), zap.Error(err))
 			}
 
-			if err := shutdowner.Shutdown(); err != nil {
+			if err2 := shutdowner.Shutdown(); err2 != nil {
+				err = errors.Join(err, err2, ErrLifecycleShutdownPanic)
 				logger.Error("encountered error trying to shutdown app: ", zap.Error(err))
 			}
 		}()

cmd/xmidt-agent/main_test.go

@@ -5,6 +5,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
@@ -14,6 +15,8 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/xmidt-org/sallust"
+	"github.com/xmidt-org/xmidt-agent/internal/websocket"
+	"go.uber.org/fx"
 )
 
 func Test_provideCLI(t *testing.T) {
@@ -132,7 +135,92 @@ func Test_xmidtAgent(t *testing.T) {
 			stopCtx, cancel := context.WithTimeout(context.Background(), time.Second)
 			defer cancel()
 			err = app.Stop(stopCtx)
-			require.NoError(err)
+			// assert.NoError(err)
+		})
+	}
+}
+
+type badShutdown struct{}
+
+func (*badShutdown) Shutdown(...fx.ShutdownOption) error {
+	return errors.New("random shutdown error")
+}
+
+func Test_xmidtAgent_rollBack(t *testing.T) {
+
+	tests := []struct {
+		description      string
+		lifeCycleOptions fx.Option
+		expectedErrs     []error
+	}{
+		// success cases
+		{
+			description: "panic in fx's Start triggered a successful rollback",
+			lifeCycleOptions: fx.Invoke(
+				func(LC fx.Lifecycle, ws *websocket.Websocket) {
+					LC.Append(
+						fx.Hook{
+							// a nil *qos.Handler will trigger the panic during fx's Start,
+							// triggering the rollback
+							OnStart: onStart(nil, ws, nil, 0, sallust.Default()),
+						},
+					)
+				},
+			),
+			expectedErrs: []error{ErrLifecycleStartPanic},
+		},
+		{
+			description: "panic in fx's Stop triggered a successful manual shutdown",
+			lifeCycleOptions: fx.Invoke(
+				func(LC fx.Lifecycle, shutdowner fx.Shutdowner) {
+					LC.Append(
+						fx.Hook{
+							// &websocket.Websocket{} will trigger the panic during fx's Stop, manually triggering
+							// the shutdown of the application by sending a signal to all open Done channels
+							OnStop: onStop(&websocket.Websocket{}, nil, shutdowner, nil, sallust.Default()),
+						},
+					)
+				},
+			),
+			expectedErrs: []error{ErrLifecycleStopPanic},
+		},
+		// fail cases
+		{
+			description: "shutdown triggered and failed",
+			lifeCycleOptions: fx.Invoke(
+				func(LC fx.Lifecycle) {
+					LC.Append(
+						fx.Hook{
+							// &websocket.Websocket{} will trigger the panic during fx's Stop, manually triggering
+							// the shutdown of the application by sending a signal to all open Done channels
+							// &badShutdown{} will trigger the panic during fx's Stop, manually triggering
+							// the shutdown of the application by sending a signal to all open Done channels
+							OnStop: onStop(&websocket.Websocket{}, nil, &badShutdown{}, nil, sallust.Default()),
+						},
+					)
+				},
+			),
+			expectedErrs: []error{ErrLifecycleStopPanic, ErrLifecycleShutdownPanic},
+		},
+	}
+	args := []string{"-f", "xmidt_agent.yaml"}
+	for _, tc := range tests {
+		t.Run(tc.description, func(t *testing.T) {
+			assert := assert.New(t)
+			app := fx.New(provideAppOptions(args), tc.lifeCycleOptions)
+
+			// only run the program for	a few seconds to make sure it starts
+			startCtx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			errs := app.Start(startCtx)
+			time.Sleep(time.Millisecond)
+
+			stopCtx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			errs = errors.Join(errs, app.Stop(stopCtx))
+			for _, err := range tc.expectedErrs {
+				assert.ErrorIs(errs, err)
+			}
 		})
 	}
 }