From 045df6b9446269495e7f7189974ec2b8f6055aef Mon Sep 17 00:00:00 2001
From: Louis Garman <75728+leg100@users.noreply.github.com>
Date: Wed, 14 Aug 2024 21:01:43 +0100
Subject: [PATCH] refactor: simplify select current workspace task (#117)

---
 internal/module/service.go    |  8 ++------
 internal/workspace/service.go | 20 ++++++--------------
 2 files changed, 8 insertions(+), 20 deletions(-)

diff --git a/internal/module/service.go b/internal/module/service.go
index 94358d28..fec8cd4d 100644
--- a/internal/module/service.go
+++ b/internal/module/service.go
@@ -251,17 +251,13 @@ func (s *Service) GetByPath(path string) (*Module, error) {
 	return nil, fmt.Errorf("%s: %w", path, resource.ErrNotFound)
 }
 
+// SetCurrent sets the current workspace for the module.
 func (s *Service) SetCurrent(moduleID, workspaceID resource.ID) error {
 	_, err := s.table.Update(moduleID, func(existing *Module) error {
 		existing.CurrentWorkspaceID = &workspaceID
 		return nil
 	})
-	if err != nil {
-		s.logger.Error("setting current workspace", "module", moduleID, "workspace", workspaceID, "error", err)
-		return err
-	}
-	s.logger.Debug("set current workspace", "module", moduleID, "workspace", workspaceID)
-	return nil
+	return err
 }
 
 // updateSpec updates the task spec with common module settings.
diff --git a/internal/workspace/service.go b/internal/workspace/service.go
index 92ba6bf5..e3a3534c 100644
--- a/internal/workspace/service.go
+++ b/internal/workspace/service.go
@@ -259,15 +259,6 @@ func (s *Service) List(opts ListOptions) []*Workspace {
 // command, which sets the current workspace for the module. Once that's
 // finished it then updates the current workspace in pug itself too.
 func (s *Service) SelectWorkspace(moduleID, workspaceID resource.ID) error {
-	if err := s.selectWorkspace(moduleID, workspaceID); err != nil {
-		s.logger.Error("selecting current workspace", "workspace_id", workspaceID, "error", err)
-		return err
-	}
-	s.logger.Debug("selected current workspace", "workspace", workspaceID)
-	return nil
-}
-
-func (s *Service) selectWorkspace(moduleID, workspaceID resource.ID) error {
 	ws, err := s.table.Get(workspaceID)
 	if err != nil {
 		return err
@@ -284,15 +275,16 @@ func (s *Service) selectWorkspace(moduleID, workspaceID resource.ID) error {
 		Args:      []string{ws.Name},
 		Immediate: true,
 		Wait:      true,
+		BeforeExited: func(t *task.Task) (task.Summary, error) {
+			// Now the terraform command has finished, update the current
+			// workspace in pug as well.
+			err := s.modules.SetCurrent(moduleID, workspaceID)
+			return nil, err
+		},
 	})
 	if err != nil {
 		return err
 	}
-	// Now task has finished successfully, update the current workspace in pug
-	// as well.
-	if err := s.modules.SetCurrent(moduleID, workspaceID); err != nil {
-		return err
-	}
 	return nil
 }