From 788af17aad5c31a06dd22716ca2a22129ccf5822 Mon Sep 17 00:00:00 2001
From: Wallace Breza <wallace.breza@microsoft.com>
Date: Thu, 18 Jan 2024 11:23:22 -0800
Subject: [PATCH] Simplified the 'WorkflowMap' type alias and unmarshalling.

---
 cli/azd/cmd/up.go                   |  2 +-
 cli/azd/pkg/workflow/config.go      | 38 +++--------------------------
 cli/azd/pkg/workflow/config_test.go | 30 +++++++++++------------
 3 files changed, 19 insertions(+), 51 deletions(-)

diff --git a/cli/azd/cmd/up.go b/cli/azd/cmd/up.go
index 18062c4f5d6..863f8243265 100644
--- a/cli/azd/cmd/up.go
+++ b/cli/azd/cmd/up.go
@@ -105,7 +105,7 @@ func (u *upAction) Run(ctx context.Context) (*actions.ActionResult, error) {
 
 	startTime := time.Now()
 
-	upWorkflow, has := u.projectConfig.Workflows.Get("up")
+	upWorkflow, has := u.projectConfig.Workflows["up"]
 	if !has {
 		upWorkflow = defaultUpWorkflow
 	} else {
diff --git a/cli/azd/pkg/workflow/config.go b/cli/azd/pkg/workflow/config.go
index 03901fe65a0..69eaa274d63 100644
--- a/cli/azd/pkg/workflow/config.go
+++ b/cli/azd/pkg/workflow/config.go
@@ -1,40 +1,7 @@
 package workflow
 
 // Stores a map of workflows configured for an azd project
-type WorkflowMap struct {
-	inner map[string]*Workflow
-}
-
-// NewWorkflowMap creates a new WorkflowMap.
-func NewWorkflowMap() *WorkflowMap {
-	return &WorkflowMap{
-		inner: map[string]*Workflow{},
-	}
-}
-
-// Set adds or updates a key-value pair in the map.
-func (wm *WorkflowMap) Set(key string, value *Workflow) {
-	if wm.inner == nil {
-		wm.inner = map[string]*Workflow{}
-	}
-
-	wm.inner[key] = value
-}
-
-// Get retrieves the value for a given key from the map.
-func (wm *WorkflowMap) Get(key string) (*Workflow, bool) {
-	if wm.inner == nil {
-		wm.inner = map[string]*Workflow{}
-	}
-
-	val, ok := wm.inner[key]
-	return val, ok
-}
-
-// MarshalYAML marshals the WorkflowMap into YAML.
-func (wm *WorkflowMap) MarshalYAML() (interface{}, error) {
-	return wm.inner, nil
-}
+type WorkflowMap map[string]*Workflow
 
 // UnmarshalYAML will unmarshal the WorkflowMap from YAML.
 // The unmarshalling will marshall the YAML like a standard Go map
@@ -49,6 +16,7 @@ func (wm *WorkflowMap) UnmarshalYAML(unmarshal func(interface{}) error) error {
 		workflow.Name = key
 	}
 
-	wm.inner = m
+	*wm = m
+
 	return nil
 }
diff --git a/cli/azd/pkg/workflow/config_test.go b/cli/azd/pkg/workflow/config_test.go
index 32e42bae71e..b748028857d 100644
--- a/cli/azd/pkg/workflow/config_test.go
+++ b/cli/azd/pkg/workflow/config_test.go
@@ -18,14 +18,14 @@ var testWorkflow = &Workflow{
 }
 
 func Test_WorkflowMap_MarshalYAML(t *testing.T) {
-	m := NewWorkflowMap()
-	m.Set("up", testWorkflow)
-	m.Set("down", testWorkflow)
+	workflowMap := WorkflowMap{}
+	workflowMap["up"] = testWorkflow
+	workflowMap["down"] = testWorkflow
 
-	expected, err := yaml.Marshal(m.inner)
+	expected, err := yaml.Marshal(workflowMap)
 	require.NoError(t, err)
 
-	actual, err := yaml.Marshal(m)
+	actual, err := yaml.Marshal(workflowMap)
 	require.NoError(t, err)
 
 	require.NoError(t, err)
@@ -34,7 +34,7 @@ func Test_WorkflowMap_MarshalYAML(t *testing.T) {
 
 func Test_WorkflowMap_UnmarshallYAML(t *testing.T) {
 	t.Run("array style", func(t *testing.T) {
-		var wm *WorkflowMap
+		var wm WorkflowMap
 		yamlString := heredoc.Doc(`
 			up:
 			  - azd: package --all
@@ -45,14 +45,14 @@ func Test_WorkflowMap_UnmarshallYAML(t *testing.T) {
 		err := yaml.Unmarshal([]byte(yamlString), &wm)
 		require.NoError(t, err)
 
-		upWorkflow, ok := wm.Get("up")
+		upWorkflow, ok := wm["up"]
 		require.True(t, ok)
 
 		assertWorkflow(t, upWorkflow)
 	})
 
 	t.Run("map style", func(t *testing.T) {
-		var wm *WorkflowMap
+		var workflowMap WorkflowMap
 		yamlString := heredoc.Doc(`
 			up:
 			  steps:
@@ -61,17 +61,17 @@ func Test_WorkflowMap_UnmarshallYAML(t *testing.T) {
 			    - azd: deploy --all
 		`)
 
-		err := yaml.Unmarshal([]byte(yamlString), &wm)
+		err := yaml.Unmarshal([]byte(yamlString), &workflowMap)
 		require.NoError(t, err)
 
-		upWorkflow, ok := wm.Get("up")
+		upWorkflow, ok := workflowMap["up"]
 		require.True(t, ok)
 
 		assertWorkflow(t, upWorkflow)
 	})
 
 	t.Run("verbose style", func(t *testing.T) {
-		var wm *WorkflowMap
+		var workflowMap WorkflowMap
 		yamlString := heredoc.Doc(`
 			up:
 			  steps:
@@ -87,22 +87,22 @@ func Test_WorkflowMap_UnmarshallYAML(t *testing.T) {
 			          - --all
 		`)
 
-		err := yaml.Unmarshal([]byte(yamlString), &wm)
+		err := yaml.Unmarshal([]byte(yamlString), &workflowMap)
 		require.NoError(t, err)
 
-		upWorkflow, ok := wm.Get("up")
+		upWorkflow, ok := workflowMap["up"]
 		require.True(t, ok)
 
 		assertWorkflow(t, upWorkflow)
 	})
 
 	t.Run("invalid workflow", func(t *testing.T) {
-		var wm *WorkflowMap
+		var workflowMap WorkflowMap
 		yamlString := heredoc.Doc(`
 			up: provision && deploy --all && package --all
 		`)
 
-		err := yaml.Unmarshal([]byte(yamlString), &wm)
+		err := yaml.Unmarshal([]byte(yamlString), &workflowMap)
 		require.Error(t, err)
 	})
 }