From a52f13fd75a70c434f93672ca575017a4a1d3fed Mon Sep 17 00:00:00 2001
From: kushalmalani <kushal@astronomer.io>
Date: Tue, 13 Jun 2023 08:48:13 -0700
Subject: [PATCH] fix refresh tokens (#1253)

* fix refresh tokens

* fix context key case

* add more context keys
---
 cmd/cloud/migration.go      | 41 --------------------
 cmd/cloud/migration_test.go | 75 -------------------------------------
 cmd/cloud/setup.go          | 17 +++++++--
 3 files changed, 13 insertions(+), 120 deletions(-)
 delete mode 100644 cmd/cloud/migration.go
 delete mode 100644 cmd/cloud/migration_test.go

diff --git a/cmd/cloud/migration.go b/cmd/cloud/migration.go
deleted file mode 100644
index a13bfa341..000000000
--- a/cmd/cloud/migration.go
+++ /dev/null
@@ -1,41 +0,0 @@
-package cloud
-
-import (
-	"errors"
-
-	astrocore "github.com/astronomer/astro-cli/astro-client-core"
-	"github.com/astronomer/astro-cli/cloud/organization"
-	"github.com/astronomer/astro-cli/context"
-)
-
-var ErrorShortName = errors.New("cannot find organization shortname")
-
-// migrate config
-func migrateCloudConfig(coreClient astrocore.CoreClient) error {
-	ctx, err := context.GetCurrentContext()
-	if err != nil {
-		return err
-	}
-	// backfill OrganizationShortName
-	if (ctx.OrganizationShortName == "") && (ctx.Organization != "") {
-		orgs, err := organization.ListOrganizations(coreClient)
-		if err != nil {
-			return err
-		}
-		shortName := ""
-		for i := range orgs {
-			if orgs[i].Id == ctx.Organization {
-				shortName = orgs[i].ShortName
-				break
-			}
-		}
-		if shortName == "" {
-			return ErrorShortName
-		}
-		err = ctx.SetContextKey("organization_short_name", shortName)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
diff --git a/cmd/cloud/migration_test.go b/cmd/cloud/migration_test.go
deleted file mode 100644
index 58ce81e1b..000000000
--- a/cmd/cloud/migration_test.go
+++ /dev/null
@@ -1,75 +0,0 @@
-package cloud
-
-import (
-	"encoding/json"
-	"net/http"
-	"testing"
-
-	astrocore "github.com/astronomer/astro-cli/astro-client-core"
-	astrocore_mocks "github.com/astronomer/astro-cli/astro-client-core/mocks"
-	"github.com/astronomer/astro-cli/context"
-	testUtil "github.com/astronomer/astro-cli/pkg/testing"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-)
-
-func TestMigrateOrgShortName(t *testing.T) {
-	var (
-		mockOKResponse = astrocore.ListOrganizationsResponse{
-			HTTPResponse: &http.Response{
-				StatusCode: 200,
-			},
-			JSON200: &[]astrocore.Organization{
-				{AuthServiceId: "auth-service-id", Id: "test-org-id", Name: "test org", ShortName: "testorg"},
-				{AuthServiceId: "auth-service-id-2", Id: "test-org-id-2", Name: "test org 2", ShortName: "testorg2"},
-			},
-		}
-		errorBody, _ = json.Marshal(astrocore.Error{
-			Message: "failed to fetch organizations",
-		})
-		mockErrorResponse = astrocore.ListOrganizationsResponse{
-			HTTPResponse: &http.Response{
-				StatusCode: 500,
-			},
-			Body:    errorBody,
-			JSON200: nil,
-		}
-	)
-	mockClient := new(astrocore_mocks.ClientWithResponsesInterface)
-
-	t.Run("can successfully backfill org shortname", func(t *testing.T) {
-		testUtil.InitTestConfig(testUtil.CloudPlatform)
-		ctx, err := context.GetCurrentContext()
-		assert.NoError(t, err)
-		ctx.SetContextKey("organization", "test-org-id")
-		ctx.SetContextKey("organization_short_name", "")
-		mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astrocore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once()
-		err = migrateCloudConfig(mockClient)
-		assert.NoError(t, err)
-		ctx, err = context.GetCurrentContext()
-		assert.NoError(t, err)
-		assert.Equal(t, "testorg", ctx.OrganizationShortName)
-	})
-
-	t.Run("throw error when ListOrganizations failed", func(t *testing.T) {
-		testUtil.InitTestConfig(testUtil.CloudPlatform)
-		ctx, err := context.GetCurrentContext()
-		assert.NoError(t, err)
-		ctx.SetContextKey("organization", "test-org-id")
-		ctx.SetContextKey("organization_short_name", "")
-		mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astrocore.ListOrganizationsParams{}).Return(&mockErrorResponse, nil).Once()
-		err = migrateCloudConfig(mockClient)
-		assert.Contains(t, err.Error(), "failed to fetch organizations")
-	})
-
-	t.Run("throw error when no organization matched", func(t *testing.T) {
-		testUtil.InitTestConfig(testUtil.CloudPlatform)
-		ctx, err := context.GetCurrentContext()
-		assert.NoError(t, err)
-		ctx.SetContextKey("organization", "test-org-id3")
-		ctx.SetContextKey("organization_short_name", "")
-		mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astrocore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once()
-		err = migrateCloudConfig(mockClient)
-		assert.Contains(t, err.Error(), "cannot find organization shortname")
-	})
-}
diff --git a/cmd/cloud/setup.go b/cmd/cloud/setup.go
index a65131e31..f1205f40f 100644
--- a/cmd/cloud/setup.go
+++ b/cmd/cloud/setup.go
@@ -125,10 +125,7 @@ func Setup(cmd *cobra.Command, client astro.Client, coreClient astrocore.CoreCli
 	if err != nil {
 		return err
 	}
-	err = migrateCloudConfig(coreClient)
-	if err != nil {
-		return err
-	}
+
 	return nil
 }
 
@@ -177,6 +174,18 @@ func checkToken(client astro.Client, coreClient astrocore.CoreClient, out io.Wri
 		if err != nil {
 			return err
 		}
+		err = c.SetContextKey("organization", c.Organization)
+		if err != nil {
+			return err
+		}
+		err = c.SetContextKey("organization_short_name", c.OrganizationShortName)
+		if err != nil {
+			return err
+		}
+		err = c.SetContextKey("organization_product", c.OrganizationProduct)
+		if err != nil {
+			return err
+		}
 	}
 	return nil
 }