From c42d666061af053f3f8890106f0a8d7ba78b8eef Mon Sep 17 00:00:00 2001
From: Keith Zantow <kzantow@gmail.com>
Date: Wed, 23 Aug 2023 13:26:30 -0400
Subject: [PATCH] chore: add post run hook (#23)

Signed-off-by: Keith Zantow <kzantow@gmail.com>
---
 application.go      | 27 +++++++++++++++++++++++++--
 application_test.go | 28 ++++++++++++++++++++++++++++
 setup_config.go     |  6 ++++++
 3 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/application.go b/application.go
index 5ac45d2..aaab8c5 100644
--- a/application.go
+++ b/application.go
@@ -21,6 +21,8 @@ import (
 
 type Initializer func(*State) error
 
+type PostRun func(*State, error)
+
 type postConstruct func(*application)
 
 type Application interface {
@@ -112,16 +114,37 @@ func (a *application) runInitializers() error {
 	return nil
 }
 
+func (a *application) runPostRuns(err error) {
+	for _, postRun := range a.setupConfig.postRuns {
+		a.runPostRun(postRun, err)
+	}
+}
+
+func (a *application) runPostRun(fn PostRun, err error) {
+	defer func() {
+		// handle panics in each postRun -- the app may already be in a panicking situation,
+		// this recover should not affect the original panic, as it is being run in a
+		// different call stack from the original panic, but the original panic should
+		// return without confusing things when a postRun also fails by panic
+		if v := recover(); v != nil {
+			a.state.Logger.Debugf("panic while calling postRun: %v", v)
+		}
+	}()
+	fn(&a.state, err)
+}
+
 func (a *application) WrapRunE(fn func(cmd *cobra.Command, args []string) error) func(cmd *cobra.Command, args []string) error {
 	return func(cmd *cobra.Command, args []string) error {
-		wrapper := func(cmd *cobra.Command, args []string) error {
+		wrapper := func(cmd *cobra.Command, args []string) (err error) {
 			defer func() {
 				// when the worker has completed (or errored) we want to exit the event loop gracefully
 				if a.state.Bus != nil {
 					a.state.Bus.Publish(ExitEvent(false))
 				}
 			}()
-			return fn(cmd, args)
+			defer a.runPostRuns(err)
+			err = fn(cmd, args)
+			return
 		}
 
 		return a.execute(cmd.Context(), async(cmd, args, wrapper))
diff --git a/application_test.go b/application_test.go
index 53ff95a..a97967d 100644
--- a/application_test.go
+++ b/application_test.go
@@ -242,6 +242,34 @@ func Test_Application_Setup_RunsInitializers(t *testing.T) {
 	require.NoError(t, cmd.Execute())
 }
 
+func Test_Application_Setup_RunsPostRuns(t *testing.T) {
+	name := "puppy"
+	version := "2.0"
+
+	cfg := NewSetupConfig(Identification{Name: name, Version: version}).WithPostRuns(
+		func(state *State, err error) {
+			t.Setenv("PUPPY_THING_STUFF", "bark-bark!")
+		},
+	)
+
+	t.Setenv("PUPPY_THING_STUFF", "ruff-ruff!")
+
+	app := New(*cfg)
+
+	app.SetupRootCommand(
+		&cobra.Command{
+			DisableFlagParsing: true,
+			Args:               cobra.ArbitraryArgs,
+			RunE: func(cmd *cobra.Command, args []string) error {
+				return nil
+			},
+		})
+
+	app.Run()
+
+	assert.Equal(t, "bark-bark!", os.Getenv("PUPPY_THING_STUFF"))
+}
+
 func Test_SetupCommand(t *testing.T) {
 	p := &persistent{}
 
diff --git a/setup_config.go b/setup_config.go
index 05ff77c..a39bf93 100644
--- a/setup_config.go
+++ b/setup_config.go
@@ -24,6 +24,7 @@ type SetupConfig struct {
 	UIConstructor     UIConstructor
 	Initializers      []Initializer
 	postConstructs    []postConstruct
+	postRuns          []PostRun
 }
 
 func NewSetupConfig(id Identification) *SetupConfig {
@@ -97,6 +98,11 @@ func (c *SetupConfig) WithInitializers(initializers ...Initializer) *SetupConfig
 	return c
 }
 
+func (c *SetupConfig) WithPostRuns(postRuns ...PostRun) *SetupConfig {
+	c.postRuns = append(c.postRuns, postRuns...)
+	return c
+}
+
 func (c *SetupConfig) withPostConstructs(postConstructs ...postConstruct) *SetupConfig {
 	c.postConstructs = append(c.postConstructs, postConstructs...)
 	return c