Add ability to switch workspaces with name (#1655)

* Add ability to switch workspaces with name * Fixing lint and test errors * Updating deployment inspect text * Update error message
astronomer · May 22, 2024 · d858097 · d858097
1 parent 63fc28c
commit d858097
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 24 deletions.
diff --git a/cloud/workspace/workspace.go b/cloud/workspace/workspace.go
@@ -120,33 +120,42 @@ var GetWorkspaceSelection = func(client astrocore.CoreClient, out io.Writer) (st
 	return selected.Id, nil
 }
 
-func Switch(id string, client astrocore.CoreClient, out io.Writer) error {
-	if id == "" {
-		_id, err := GetWorkspaceSelection(client, out)
+func Switch(workspaceNameOrID string, client astrocore.CoreClient, out io.Writer) error {
+	var wsID string
+	if workspaceNameOrID == "" {
+		id, err := GetWorkspaceSelection(client, out)
+		if err != nil {
+			return err
+		}
+
+		wsID = id
+	} else {
+		ws, err := GetWorkspaces(client)
 		if err != nil {
 			return err
 		}
+		for i := range ws {
+			if ws[i].Name == workspaceNameOrID || ws[i].Id == workspaceNameOrID {
+				wsID = ws[i].Id
+			}
+		}
 
-		id = _id
+		if wsID == "" {
+			return errors.Wrap(err, "workspace id/name could not be found")
+		}
 	}
 
 	c, err := config.GetCurrentContext()
 	if err != nil {
 		return err
 	}
 
-	// validate workspace
-	_, err = GetWorkspaces(client)
-	if err != nil {
-		return errors.Wrap(err, "workspace id is not valid")
-	}
-
-	err = c.SetContextKey("workspace", id)
+	err = c.SetContextKey("workspace", wsID)
 	if err != nil {
 		return err
 	}
 
-	err = c.SetContextKey("last_used_workspace", id)
+	err = c.SetContextKey("last_used_workspace", wsID)
 	if err != nil {
 		return err
 	}

diff --git a/cloud/workspace/workspace_test.go b/cloud/workspace/workspace_test.go
@@ -146,27 +146,51 @@ func TestSwitch(t *testing.T) {
 	testUtil.InitTestConfig(testUtil.LocalPlatform)
 	mockCoreClient := new(astrocore_mocks.ClientWithResponsesInterface)
 
+	workspace2 := astrocore.Workspace{
+		Name:                         "test-workspace-2",
+		Description:                  &description,
+		ApiKeyOnlyDeploymentsDefault: false,
+		Id:                           "workspace-id-2",
+	}
+
+	workspaces = []astrocore.Workspace{
+		workspace1,
+		workspace2,
+	}
+
+	ListWorkspacesResponseOK = astrocore.ListWorkspacesResponse{
+		HTTPResponse: &http.Response{
+			StatusCode: 200,
+		},
+		JSON200: &astrocore.WorkspacesPaginated{
+			Limit:      10,
+			Offset:     0,
+			TotalCount: 2,
+			Workspaces: workspaces,
+		},
+	}
+
 	t.Run("success", func(t *testing.T) {
 		mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once()
 
 		buf := new(bytes.Buffer)
-		err := Switch("test-id-1", mockCoreClient, buf)
+		err := Switch("workspace-id-2", mockCoreClient, buf)
 		assert.NoError(t, err)
-		assert.Contains(t, buf.String(), "test-id-1")
+		assert.Contains(t, buf.String(), "workspace-id-2")
 		mockCoreClient.AssertExpectations(t)
 	})
 
 	t.Run("list workspace failure", func(t *testing.T) {
 		mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(nil, errMock).Once()
 
 		buf := new(bytes.Buffer)
-		err := Switch("test-id-1", mockCoreClient, buf)
+		err := Switch("workspace-id-2", mockCoreClient, buf)
 		assert.ErrorIs(t, err, errMock)
 		mockCoreClient.AssertExpectations(t)
 	})
 
 	t.Run("success with selection", func(t *testing.T) {
-		mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Twice()
+		mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once()
 
 		// mock os.Stdin
 		input := []byte("1")

diff --git a/cmd/cloud/deployment_inspect.go b/cmd/cloud/deployment_inspect.go
@@ -19,8 +19,8 @@ func newDeploymentInspectCmd(out io.Writer) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:     "inspect",
 		Aliases: []string{"in"},
-		Short:   "Inspect a deployment",
-		Long:    "Inspect an Astro Deployment.",
+		Short:   "Inspect a deployment configuration",
+		Long:    "Inspect an Astro Deployment configuration, which can be useful if you manage deployments as code or use Deployment configuration templates. This command returns the Deployment's configuration as a YAML or JSON output, which includes information about resources, such as cluster ID, region, and Airflow API URL, as well as scheduler and worker queue configurations.",
 		RunE: func(cmd *cobra.Command, args []string) error {
 			return deploymentInspect(cmd, args, out)
 		},

diff --git a/cmd/cloud/deployment_inspect_test.go b/cmd/cloud/deployment_inspect_test.go
@@ -11,7 +11,7 @@ import (
 )
 
 func TestNewDeploymentInspectCmd(t *testing.T) {
-	expectedHelp := "Inspect an Astro Deployment."
+	expectedHelp := "Inspect an Astro Deployment configuration, which can be useful if you manage deployments as code or use Deployment configuration templates. This command returns the Deployment's configuration as a YAML or JSON output, which includes information about resources, such as cluster ID, region, and Airflow API URL, as well as scheduler and worker queue configurations."
 	testUtil.InitTestConfig(testUtil.LocalPlatform)
 	mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface)
 	platformCoreClient = mockPlatformCoreClient

diff --git a/cmd/cloud/workspace.go b/cmd/cloud/workspace.go
@@ -85,7 +85,7 @@ func newWorkspaceListCmd(out io.Writer) *cobra.Command {
 
 func newWorkspaceSwitchCmd(out io.Writer) *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "switch [workspace_id]",
+		Use:     "switch [workspace name/id]",
 		Aliases: []string{"sw"},
 		Short:   "Switch to a different Astro Workspace",
 		Long:    "Switch to a different Astro Workspace",
@@ -598,13 +598,13 @@ func workspaceList(cmd *cobra.Command, out io.Writer) error {
 func workspaceSwitch(cmd *cobra.Command, out io.Writer, args []string) error {
 	// Silence Usage as we have now validated command input
 
-	id := ""
+	workspaceNameOrID := ""
 
 	if len(args) == 1 {
-		id = args[0]
+		workspaceNameOrID = args[0]
 	}
 	cmd.SilenceUsage = true
-	return workspace.Switch(id, astroCoreClient, out)
+	return workspace.Switch(workspaceNameOrID, astroCoreClient, out)
 }
 
 func workspaceCreate(cmd *cobra.Command, out io.Writer) error {

diff --git a/cmd/cloud/workspace_test.go b/cmd/cloud/workspace_test.go
@@ -86,7 +86,7 @@ func TestWorkspaceSwitch(t *testing.T) {
 	testUtil.InitTestConfig(testUtil.LocalPlatform)
 
 	mockClient := new(astrocore_mocks.ClientWithResponsesInterface)
-	mockClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Twice()
+	mockClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once()
 	astroCoreClient = mockClient
 
 	// mock os.Stdin