when there is more then one org for the auth key

open Quick-pick dialog and let the user pick the org

.eslintrc.json

@@ -20,6 +20,7 @@
     ],
     "rules": {
         "no-useless-constructor": "off",
-        "@typescript-eslint/ban-ts-comment": "off"
+        "@typescript-eslint/ban-ts-comment": "off",
+        "@typescript-eslint/no-non-null-assertion": "off"
       }
 }

src/api-client.ts

@@ -5,6 +5,7 @@ import { DeploymentStepLogsResponse, DeploymentStepResponse } from "./types";
 
 class ApiClient {
   private credentials?: { username: string; password: string };
+  private currentOrganizationId?: string;
   private readonly instance: AxiosInstance;
   constructor() {
     this.instance = axios.create({ baseURL: `https://${ENV0_API_URL}` });
@@ -16,11 +17,16 @@ class ApiClient {
     });
   }
 
-  public init(credentials: { username: string; password: string }) {
+  public init(
+    credentials: { username: string; password: string },
+    organizationId: string
+  ) {
     this.credentials = credentials;
+    this.currentOrganizationId = organizationId;
   }
 
   public clearCredentials() {
+    this.currentOrganizationId = undefined;
     this.credentials = undefined;
   }
 
@@ -62,11 +68,11 @@ class ApiClient {
     return response.data;
   }
 
-  public async getEnvironments(organizationId: string) {
+  public async getEnvironments() {
     const response = await this.instance.get<EnvironmentModel[]>(
       `/environments`,
       {
-        params: { organizationId, isActive: true },
+        params: { organizationId: this.currentOrganizationId, isActive: true },
       }
     );
 

src/auth.ts

@@ -11,6 +11,13 @@ const env0SecretKey = "env0.secret";
 
 export class AuthService {
   constructor(private readonly context: vscode.ExtensionContext) {}
+
+  private selectedOrgId: string | undefined;
+
+  public getSelectedOrg() {
+    return this.selectedOrgId;
+  }
+
   public registerLoginCommand(onLogin: () => void) {
     const disposable = vscode.commands.registerCommand(
       "env0.login",
@@ -47,9 +54,9 @@ export class AuthService {
           },
         });
 
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        if (await this.validateUserCredentials(keyId!, secret!)) {
-          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        const selectedOrgId = await this.pickOrganization(keyId!, secret!);
+        if (selectedOrgId) {
+          this.selectedOrgId = selectedOrgId;
           await this.storeAuthData(keyId!, secret!);
           await onLogin();
         }
@@ -89,27 +96,50 @@ export class AuthService {
     };
   }
 
-  private async validateUserCredentials(keyId: string, secret: string) {
+  private async pickOrganization(keyId: string, secret: string) {
     // Displaying a loading indicator to inform the user that something is happening
     return await vscode.window.withProgress(
       { location: { viewId: ENV0_ENVIRONMENTS_VIEW_ID } },
       async () => {
         try {
-          await axios.get(`https://${ENV0_API_URL}/organizations`, {
-            auth: { username: keyId, password: secret },
-            validateStatus: function (status) {
-              return status >= 200 && status < 300;
-            },
+          const orgsRes = await axios.get(
+            `https://${ENV0_API_URL}/organizations`,
+            {
+              auth: { username: keyId, password: secret },
+              validateStatus: function (status) {
+                return status >= 200 && status < 300;
+              },
+            }
+          );
+          if (orgsRes.data.length === 1) {
+            return orgsRes.data[0].id;
+          }
+          const orgs = orgsRes.data.map((org: any) => ({
+            name: org.name,
+            id: org.id,
+          }));
+          const items: vscode.QuickPickItem[] = orgs.map((org: any) => ({
+            label: org.name,
+            description: org.id,
+          }));
+
+          const selectedItem = await vscode.window.showQuickPick(items, {
+            placeHolder: "Select an organization",
           });
-          return true;
+          const selectedOrgId = selectedItem?.description;
+          if (!selectedOrgId) {
+            vscode.window.showErrorMessage("No organization selected");
+            return undefined;
+          }
+          return selectedOrgId;
           // eslint-disable-next-line @typescript-eslint/no-explicit-any
         } catch (e: any) {
           if (e?.response?.status >= 400 && e?.response?.status < 500) {
             showInvalidCredentialsMessage();
           } else {
             showUnexpectedErrorMessage();
           }
-          return false;
+          return undefined;
         }
       }
     );

src/extension.ts

@@ -70,7 +70,10 @@ const init = async (
   environmentsTree: vscode.TreeView<Environment>,
   authService: AuthService
 ) => {
-  apiClient.init(await authService.getApiKeyCredentials());
+  apiClient.init(
+    await authService.getApiKeyCredentials(),
+    authService.getSelectedOrg()!
+  );
   extensionState.setLoggedIn(true);
   await loadEnvironments(environmentsDataProvider);
 };

src/get-environments.ts

@@ -38,11 +38,7 @@ function repositoriesEqual(rep1: string, rep2: string): boolean {
 export async function getEnvironmentsForBranch() {
   let environments: EnvironmentModel[] = [];
 
-  const organizationId = await getOrganizationId();
-
-  if (organizationId) {
-    environments = await getEnvironments(organizationId);
-  }
+  environments = await getEnvironments();
 
   if (environments.length > 0) {
     const { currentBranch, repository } = getGitRepoAndBranch();
@@ -64,17 +60,12 @@ export async function getEnvironmentsForBranch() {
   return environments;
 }
 
-async function getEnvironments(organizationId: string) {
+async function getEnvironments() {
   try {
-    return apiClient.getEnvironments(organizationId);
+    return apiClient.getEnvironments();
   } catch (e) {
     console.log(e);
   }
 
   return [];
 }
-
-async function getOrganizationId() {
-  const organizations = await apiClient.getOrganizations();
-  return organizations[0]?.id;
-}

src/test/integration/mocks/quick-pick.ts

@@ -0,0 +1,12 @@
+import * as jestMock from "jest-mock";
+import * as vscode from "vscode";
+
+let quickPickSpy: jestMock.SpyInstance<any>;
+export const spyOnQuickPick = () => {
+  quickPickSpy = jestMock.spyOn(vscode.window, "showQuickPick");
+  return quickPickSpy;
+};
+
+export const resetSpyOnQuickPick = () => {
+  quickPickSpy?.mockReset?.();
+};

src/test/integration/mocks/server.ts

@@ -29,15 +29,15 @@ const assertAuth = (credentials: Credentials, authHeader: string | null) => {
 };
 
 export const mockGetOrganization = (
-  organizationId: string,
+  orgs: { id: string; name: string }[],
   credentials?: Credentials
 ) => {
   server.use(
     rest.get(`https://${ENV0_API_URL}/organizations`, (req, res, ctx) => {
       if (credentials) {
         assertAuth(credentials, req.headers.get("Authorization"));
       }
-      return res(ctx.json([{ id: organizationId }]));
+      return res(ctx.json(orgs));
     })
   );
 };

src/test/integration/suite/authentication.test.it.ts

@@ -18,9 +18,9 @@ import {
 import * as jestMock from "jest-mock";
 import * as vscode from "vscode";
 import expect from "expect";
-import { afterEach, beforeEach } from "mocha";
+import { afterEach } from "mocha";
+import { resetSpyOnQuickPick, spyOnQuickPick } from "../mocks/quick-pick";
 
-const orgId = "org-id";
 const envName = "my env";
 const auth = { keyId: "key-id", secret: "key-secret" };
 const environmentMock = getEnvironmentMock(
@@ -30,40 +30,44 @@ const environmentMock = getEnvironmentMock(
     name: envName,
   }
 );
+const selectedOrg = { name: "my org", id: "org-id" };
+const initMocksAndLogin = async (moreOrgs: typeof selectedOrg[] = []) => {
+  mockGetOrganization([selectedOrg, ...moreOrgs], auth);
+  mockGetEnvironment(selectedOrg.id, [environmentMock], auth);
+  mockGitRepoAndBranch("main", "[email protected]:user/repo.git");
+  mockGetDeploymentStepsApiResponse();
+  await login(auth);
 
-suite("authentication", function () {
-  this.timeout(1000 * 10);
-  beforeEach(async () => {
-    mockGetOrganization(orgId, auth);
-    mockGetEnvironment(orgId, [environmentMock], auth);
-    mockGitRepoAndBranch("main", "[email protected]:user/repo.git");
-    mockGetDeploymentStepsApiResponse();
+  await waitFor(() => expect(getFirstEnvStatus()).toContain("ACTIVE"));
+};
 
-    await login(auth);
-    await waitFor(() => expect(getFirstEnvStatus()).toContain("ACTIVE"));
-  });
+suite("authentication", function () {
+  this.timeout(1000 * 600);
 
   afterEach(async () => {
     await logout();
     await resetExtension();
+    resetSpyOnQuickPick();
   });
 
   test("should call redeploy with the credentials provided on login", async () => {
+    await initMocksAndLogin();
     const onRedeployCalled = jestMock.fn();
     mockRedeployApiResponse(environmentMock.id, auth, onRedeployCalled);
 
     vscode.commands.executeCommand("env0.redeploy", getFirstEnvironment());
-    await waitFor(() => expect(onRedeployCalled).toHaveBeenCalled());
+    await waitFor(() => expect(onRedeployCalled).toHaveBeenCalled(), 10);
   });
 
   test("should call redeploy with updated credentials when logout and login again ", async () => {
+    await initMocksAndLogin();
     await logout();
     const newAuthData = {
       keyId: "different-key-id",
       secret: "different-key-secret",
     };
-    mockGetOrganization(orgId, newAuthData);
-    mockGetEnvironment(orgId, [environmentMock], newAuthData);
+    mockGetOrganization([selectedOrg], newAuthData);
+    mockGetEnvironment(selectedOrg.id, [environmentMock], newAuthData);
 
     await login(newAuthData);
     await waitFor(() => expect(getFirstEnvStatus()).toContain("ACTIVE"));
@@ -75,11 +79,30 @@ suite("authentication", function () {
   });
 
   test("should show login message when logout", async () => {
+    await initMocksAndLogin();
     await logout();
     await waitFor(() =>
       expect(getEnvironmentsView().message).toContain(
         "you are logged out. in order to log in, run the command 'env0.login'"
       )
     );
   });
+
+  test("should show pick organization message when login", async () => {
+    const onQuickPick = spyOnQuickPick();
+
+    const secondOrg = { name: "second org", id: "second-org-id" };
+    initMocksAndLogin([secondOrg]);
+    await waitFor(() =>
+      expect(onQuickPick).toHaveBeenCalledWith(
+        [selectedOrg, secondOrg].map((org) => ({
+          label: org.name,
+          description: org.id,
+        })),
+        {
+          placeHolder: "Select an organization",
+        }
+      )
+    );
+  });
 });

src/test/integration/suite/deployment-logs.test.it.ts

@@ -28,7 +28,8 @@ import expect from "expect";
 import { DeploymentStatus } from "../../../types";
 
 const auth = { keyId: "key-id", secret: "key-secret" };
-const orgId = "org-id";
+const organization = { name: "my org", id: "org-id" };
+const orgId = organization.id;
 const firstEnvName = "my env";
 const secondEnvName = "my env 2";
 
@@ -62,8 +63,8 @@ suite("deployment logs", function () {
     mockOutputChannel();
     await resetExtension(); // we need to resat because we are mocking the output channel
     const environments = [firstEnvironmentMock, secondEnvironmentMock];
-    mockGetOrganization(orgId, auth);
-    mockGetEnvironment(orgId, environments, auth);
+    mockGetOrganization([organization], auth);
+    mockGetEnvironment(organization.id, environments, auth);
     await login(auth);
     await waitFor(() => expect(getFirstEnvStatus()).toBe("ACTIVE"));
   });

src/test/integration/suite/environment-actions.test.it.ts

@@ -41,11 +41,12 @@ import {
 } from "../mocks/notification-message";
 
 const auth = { keyId: "key-id", secret: "key-secret" };
-const orgId = "org-id";
+const organization = { name: "my org", id: "org-id" };
+const orgId = organization.id;
 
 const initTest = async (environments: EnvironmentModel[]) => {
-  mockGetOrganization(orgId, auth);
-  mockGetEnvironment(orgId, environments, auth);
+  mockGetOrganization([organization], auth);
+  mockGetEnvironment(organization.id, environments, auth);
   mockGitRepoAndBranch("main", "[email protected]:user/repo.git");
   mockGetDeploymentStepsApiResponse();
   spyOnShowMessage();

src/test/integration/suite/environments.test.it.ts

@@ -16,10 +16,11 @@ import { EnvironmentModel } from "../../../get-environments";
 import expect from "expect";
 
 const auth = { keyId: "key-id", secret: "key-secret" };
-const orgId = "org-id";
+const organization = { name: "my org", id: "org-id" };
+
 const initTest = async (environments: EnvironmentModel[]) => {
-  mockGetOrganization(orgId, auth);
-  mockGetEnvironment(orgId, environments, auth);
+  mockGetOrganization([organization], auth);
+  mockGetEnvironment(organization.id, environments, auth);
   mockGitRepoAndBranch("main", "[email protected]:user/repo.git");
   await login(auth);
 };
@@ -80,15 +81,15 @@ suite("environments", function () {
       }),
     ];
     mockNoGitRepo();
-    mockGetOrganization(orgId, auth);
+    mockGetOrganization([organization], auth);
     // we don't await on login because we want to test the loading message
     login(auth);
     await waitFor(
       () => expect(getEnvironmentViewMessage()).toBe("loading environments..."),
       10
     );
     mockGitRepoAndBranch("main", "[email protected]:user/repo.git");
-    mockGetEnvironment(orgId, environments, auth, 2000);
+    mockGetEnvironment(organization.id, environments, auth, 2000);
     await waitFor(() =>
       expect(getEnvironmentViewMessage()).toBe(
         "loading environments for branch main..."
@@ -99,7 +100,7 @@ suite("environments", function () {
 
   test("should show could not find git branch message", async () => {
     mockNoGitRepo();
-    mockGetOrganization(orgId, auth);
+    mockGetOrganization([organization], auth);
     await login(auth);
     await waitFor(() =>
       expect(getEnvironmentViewMessage()).toBe(