From e4333c0433c7fd245a10cbfd3387fcacf9e9e179 Mon Sep 17 00:00:00 2001
From: Steven Masley <Emyrk@users.noreply.github.com>
Date: Mon, 24 Jun 2024 07:55:39 -1000
Subject: [PATCH] chore: 'coder login' reset cli organization context (#13646)

Cli organization context is reset on `coder login` if the organization
selected is an invalid organization.
---
 cli/login.go      |  7 +++++++
 cli/login_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/cli/login.go b/cli/login.go
index 87cfea103c..faec491270 100644
--- a/cli/login.go
+++ b/cli/login.go
@@ -336,6 +336,13 @@ func (r *RootCmd) login() *serpent.Command {
 				return xerrors.Errorf("write server url: %w", err)
 			}
 
+			// If the current organization cannot be fetched, then reset the organization context.
+			// Otherwise, organization cli commands will fail.
+			_, err = CurrentOrganization(r, inv, client)
+			if err != nil {
+				_ = config.Organization().Delete()
+			}
+
 			_, _ = fmt.Fprintf(inv.Stdout, Caret+"Welcome to Coder, %s! You're authenticated.\n", pretty.Sprint(cliui.DefaultStyles.Keyword, resp.Username))
 			return nil
 		},
diff --git a/cli/login_test.go b/cli/login_test.go
index 3cf9dc1945..dc551ccaf2 100644
--- a/cli/login_test.go
+++ b/cli/login_test.go
@@ -5,9 +5,11 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"runtime"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -304,4 +306,48 @@ func TestLogin(t *testing.T) {
 		// This **should not be equal** to the token we passed in.
 		require.NotEqual(t, client.SessionToken(), sessionFile)
 	})
+
+	// Login should reset the configured organization if the user is not a member
+	t.Run("ResetOrganization", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		coderdtest.CreateFirstUser(t, client)
+		root, cfg := clitest.New(t, "login", client.URL.String(), "--token", client.SessionToken())
+
+		notRealOrg := uuid.NewString()
+		err := cfg.Organization().Write(notRealOrg)
+		require.NoError(t, err, "write bad org to config")
+
+		err = root.Run()
+		require.NoError(t, err)
+		sessionFile, err := cfg.Session().Read()
+		require.NoError(t, err)
+		require.NotEqual(t, client.SessionToken(), sessionFile)
+
+		// Organization config should be deleted since the org does not exist
+		selected, err := cfg.Organization().Read()
+		require.ErrorIs(t, err, os.ErrNotExist)
+		require.NotEqual(t, selected, notRealOrg)
+	})
+
+	t.Run("KeepOrganizationContext", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		first := coderdtest.CreateFirstUser(t, client)
+		root, cfg := clitest.New(t, "login", client.URL.String(), "--token", client.SessionToken())
+
+		err := cfg.Organization().Write(first.OrganizationID.String())
+		require.NoError(t, err, "write bad org to config")
+
+		err = root.Run()
+		require.NoError(t, err)
+		sessionFile, err := cfg.Session().Read()
+		require.NoError(t, err)
+		require.NotEqual(t, client.SessionToken(), sessionFile)
+
+		// Organization config should be deleted since the org does not exist
+		selected, err := cfg.Organization().Read()
+		require.NoError(t, err)
+		require.Equal(t, selected, first.OrganizationID.String())
+	})
 }