Adds middleware to ensure azd environment is always available

cli/azd/cmd/actions/action_descriptor.go

@@ -148,6 +148,13 @@ type CommandGroupOptions struct {
 	RootLevelHelp RootLevelHelpOption
 }
 
+// EnvironmentOptions contains options for the environment flag and initialization
+type EnvironmentOptions struct {
+	// Optional should be set to true when an azd environment is optional within an azd command.
+	// Well known use cases for this are for `azd init` and `azd show`
+	Optional bool
+}
+
 // Defines the type used for annotating a command as part of a group.
 type commandGroupAnnotationKey string
 
@@ -186,6 +193,8 @@ type ActionDescriptorOptions struct {
 	HelpOptions ActionHelpOptions
 	// Defines grouping options for the command
 	GroupingOptions CommandGroupOptions
+	// Defines options for the construction of the azd environment
+	Environment EnvironmentOptions
 }
 
 // Completion function used for cobra command flag completion

cli/azd/cmd/container.go

@@ -195,50 +195,13 @@ func registerCommonDependencies(container *ioc.NestedContainer) {
 	// Register an initialized environment based on the specified environment flag, or the default environment.
 	// Note that referencing an *environment.Environment in a command automatically triggers a UI prompt if the
 	// environment is uninitialized or a default environment doesn't yet exist.
-	container.MustRegisterScoped(
-		func(ctx context.Context,
-			azdContext *azdcontext.AzdContext,
-			envManager environment.Manager,
-			lazyEnv *lazy.Lazy[*environment.Environment],
-			envFlags internal.EnvFlag,
-		) (*environment.Environment, error) {
-			if azdContext == nil {
-				return nil, azdcontext.ErrNoProject
-			}
-
-			environmentName := envFlags.EnvironmentName
-			var err error
-
-			env, err := envManager.LoadOrInitInteractive(ctx, environmentName)
-			if err != nil {
-				return nil, fmt.Errorf("loading environment: %w", err)
-			}
-
-			// Reset lazy env value after loading or creating environment
-			// This allows any previous lazy instances (such as hooks) to now point to the same instance
-			lazyEnv.SetValue(env)
+	container.MustRegisterScoped(func(lazyEnv *lazy.Lazy[*environment.Environment]) (*environment.Environment, error) {
+		return lazyEnv.GetValue()
+	})
 
-			return env, nil
-		},
-	)
-	container.MustRegisterScoped(func(lazyEnvManager *lazy.Lazy[environment.Manager]) environment.EnvironmentResolver {
+	container.MustRegisterScoped(func(lazyEnv *lazy.Lazy[*environment.Environment]) environment.EnvironmentResolver {
 		return func(ctx context.Context) (*environment.Environment, error) {
-			azdCtx, err := azdcontext.NewAzdContext()
-			if err != nil {
-				return nil, err
-			}
-			defaultEnv, err := azdCtx.GetDefaultEnvironmentName()
-			if err != nil {
-				return nil, err
-			}
-
-			// We need to lazy load the environment manager since it depends on azd context
-			envManager, err := lazyEnvManager.GetValue()
-			if err != nil {
-				return nil, err
-			}
-
-			return envManager.Get(ctx, defaultEnv)
+			return lazyEnv.GetValue()
 		}
 	})
 

cli/azd/cmd/middleware/environment.go

@@ -0,0 +1,82 @@
+package middleware
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/azure/azure-dev/cli/azd/cmd/actions"
+	"github.com/azure/azure-dev/cli/azd/internal"
+	"github.com/azure/azure-dev/cli/azd/pkg/environment"
+	"github.com/azure/azure-dev/cli/azd/pkg/environment/azdcontext"
+	"github.com/azure/azure-dev/cli/azd/pkg/lazy"
+)
+
+// EnvironmentMiddleware is a middleware that loads the environment when not readily available
+type EnvironmentMiddleware struct {
+	lazyAzdContext *lazy.Lazy[*azdcontext.AzdContext]
+	lazyEnvManager *lazy.Lazy[environment.Manager]
+	lazyEnv        *lazy.Lazy[*environment.Environment]
+	envFlags       internal.EnvFlag
+}
+
+// NewEnvironmentMiddleware creates a new instance of the EnvironmentMiddleware
+func NewEnvironmentMiddleware(
+	lazyAzdContext *lazy.Lazy[*azdcontext.AzdContext],
+	lazyEnvManager *lazy.Lazy[environment.Manager],
+	lazyEnv *lazy.Lazy[*environment.Environment],
+	envFlags internal.EnvFlag,
+) Middleware {
+	return &EnvironmentMiddleware{
+		lazyAzdContext: lazyAzdContext,
+		lazyEnvManager: lazyEnvManager,
+		lazyEnv:        lazyEnv,
+		envFlags:       envFlags,
+	}
+}
+
+// Run runs the EnvironmentMiddleware to load the environment when not readily available
+func (m *EnvironmentMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
+	// We already have an environment, skip loading
+	// This will typically be the case when an environment has been created from a previous command like `azd init`
+	env, err := m.lazyEnv.GetValue()
+	if err == nil && env != nil {
+		return next(ctx)
+	}
+
+	// Needs Azd context before we can have an environment
+	azdContext, err := m.lazyAzdContext.GetValue()
+	if err != nil {
+		// No Azd context errors will by handled downstream
+		return next(ctx)
+	}
+
+	envManager, err := m.lazyEnvManager.GetValue()
+	if err != nil {
+		return nil, fmt.Errorf("loading environment manager: %w", err)
+	}
+
+	// Check env flag (-e, --environment) and environment variable (AZURE_ENV_NAME)
+	environmentName := m.envFlags.EnvironmentName
+	if environmentName == "" {
+		environmentName, err = azdContext.GetDefaultEnvironmentName()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// Load or initialize environment interactively from user prompt
+	env, err = envManager.LoadOrInitInteractive(ctx, environmentName)
+	if err != nil {
+		//nolint:lll
+		return nil, fmt.Errorf(
+			"failed loading environment. Ensure environment has been set using flag (--environment, -e) or by setting environment variable 'AZURE_ENV_NAME'. %w",
+			err,
+		)
+	}
+
+	// Reset lazy env value after loading or creating environment
+	// This allows any previous lazy instances (such as hooks) to now point to the same instance
+	m.lazyEnv.SetValue(env)
+
+	return next(ctx)
+}

cli/azd/cmd/middleware/environment_test.go

@@ -0,0 +1,151 @@
+package middleware
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/azure/azure-dev/cli/azd/cmd/actions"
+	"github.com/azure/azure-dev/cli/azd/internal"
+	"github.com/azure/azure-dev/cli/azd/pkg/environment"
+	"github.com/azure/azure-dev/cli/azd/pkg/environment/azdcontext"
+	"github.com/azure/azure-dev/cli/azd/pkg/lazy"
+	"github.com/azure/azure-dev/cli/azd/test/mocks"
+	"github.com/azure/azure-dev/cli/azd/test/mocks/mockenv"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Environment_Already_Exists(t *testing.T) {
+	expectedEnv := environment.NewWithValues("test", map[string]string{
+		environment.SubscriptionIdEnvVarName: "SUBSCRIPTION_ID",
+	})
+
+	mockContext := mocks.NewMockContext(context.Background())
+	azdContext := azdcontext.NewAzdContextWithDirectory(t.TempDir())
+
+	middleware, lazyEnv := createMiddlewareForTest(azdContext, expectedEnv, internal.EnvFlag{}, &mockenv.MockEnvManager{})
+	result, err := middleware.Run(*mockContext.Context, nextFn)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	actualEnv, err := lazyEnv.GetValue()
+	require.NoError(t, err)
+	require.NotNil(t, actualEnv)
+	require.Equal(t, expectedEnv.Name(), actualEnv.Name())
+}
+
+func Test_Environment_No_Azd_Context(t *testing.T) {
+	mockContext := mocks.NewMockContext(context.Background())
+
+	middleware, lazyEnv := createMiddlewareForTest(nil, nil, internal.EnvFlag{}, &mockenv.MockEnvManager{})
+	result, err := middleware.Run(*mockContext.Context, nextFn)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	actualEnv, err := lazyEnv.GetValue()
+	require.Error(t, err)
+	require.Nil(t, actualEnv)
+}
+
+func Test_Environment_With_Flag(t *testing.T) {
+	expectedEnv := environment.NewWithValues("flag-env", map[string]string{})
+
+	mockContext := mocks.NewMockContext(context.Background())
+	azdContext := azdcontext.NewAzdContextWithDirectory(t.TempDir())
+	envFlag := internal.EnvFlag{EnvironmentName: expectedEnv.Name()}
+
+	envManager := &mockenv.MockEnvManager{}
+	envManager.On("LoadOrInitInteractive", mock.Anything, mock.Anything).Return(expectedEnv, nil)
+
+	middleware, lazyEnv := createMiddlewareForTest(azdContext, nil, envFlag, envManager)
+	result, err := middleware.Run(*mockContext.Context, nextFn)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	actualEnv, err := lazyEnv.GetValue()
+	require.NoError(t, err)
+	require.NotNil(t, actualEnv)
+	require.Equal(t, expectedEnv.Name(), actualEnv.Name())
+}
+
+func Test_Environment_From_Prompt(t *testing.T) {
+	expectedEnv := environment.NewWithValues("prompt-env", map[string]string{})
+
+	mockContext := mocks.NewMockContext(context.Background())
+	azdContext := azdcontext.NewAzdContextWithDirectory(t.TempDir())
+
+	envManager := &mockenv.MockEnvManager{}
+	envManager.On("LoadOrInitInteractive", mock.Anything, mock.Anything).Return(expectedEnv, nil)
+
+	middleware, lazyEnv := createMiddlewareForTest(azdContext, nil, internal.EnvFlag{}, envManager)
+	result, err := middleware.Run(*mockContext.Context, nextFn)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	actualEnv, err := lazyEnv.GetValue()
+	require.NoError(t, err)
+	require.NotNil(t, actualEnv)
+	require.Equal(t, expectedEnv.Name(), actualEnv.Name())
+}
+
+func Test_Environment_From_Default(t *testing.T) {
+	expectedEnv := environment.NewWithValues("default-env", map[string]string{})
+
+	mockContext := mocks.NewMockContext(context.Background())
+	azdContext := azdcontext.NewAzdContextWithDirectory(t.TempDir())
+	err := azdContext.SetProjectState(azdcontext.ProjectState{
+		DefaultEnvironment: expectedEnv.Name(),
+	})
+	require.NoError(t, err)
+
+	envManager := &mockenv.MockEnvManager{}
+	envManager.On("LoadOrInitInteractive", mock.Anything, mock.Anything).Return(expectedEnv, nil)
+
+	middleware, lazyEnv := createMiddlewareForTest(azdContext, nil, internal.EnvFlag{}, envManager)
+	result, err := middleware.Run(*mockContext.Context, nextFn)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	actualEnv, err := lazyEnv.GetValue()
+	require.NoError(t, err)
+	require.NotNil(t, actualEnv)
+	require.Equal(t, expectedEnv.Name(), actualEnv.Name())
+}
+
+func createMiddlewareForTest(
+	azdContext *azdcontext.AzdContext,
+	env *environment.Environment,
+	envFlag internal.EnvFlag,
+	mockEnvManager *mockenv.MockEnvManager,
+) (Middleware, *lazy.Lazy[*environment.Environment]) {
+	// Setup environment mocks for save & reload
+	mockEnvManager.On("Save", mock.Anything, mock.Anything).Return(nil)
+	mockEnvManager.On("Reload", mock.Anything, mock.Anything).Return(nil)
+
+	lazyAzdContext := lazy.NewLazy(func() (*azdcontext.AzdContext, error) {
+		if azdContext == nil {
+			return nil, azdcontext.ErrNoProject
+		}
+
+		return azdContext, nil
+	})
+
+	lazyEnvManager := lazy.NewLazy(func() (environment.Manager, error) {
+		return mockEnvManager, nil
+	})
+
+	lazyEnv := lazy.NewLazy(func() (*environment.Environment, error) {
+		if env == nil {
+			return nil, errors.New("environemnt not found")
+		}
+
+		return env, nil
+	})
+
+	return NewEnvironmentMiddleware(lazyAzdContext, lazyEnvManager, lazyEnv, envFlag), lazyEnv
+}
+
+func nextFn(ctx context.Context) (*actions.ActionResult, error) {
+	return &actions.ActionResult{}, nil
+}