feat: display tenant name in account tree item (#12553)

* feat: display tenant name in account tree item

* feat: add feature flag

* feat: display tenant name for azure account

* test: fix ut fail

* test: add ut

* test: add ut to raise code coverage

* test: add ut to raise code coverage

packages/fx-core/src/common/featureFlags.ts

@@ -34,6 +34,7 @@ export class FeatureFlagName {
   static readonly KiotaIntegration = "TEAMSFX_KIOTA_INTEGRATION";
   static readonly ApiPluginAAD = "TEAMSFX_API_PLUGIN_AAD";
   static readonly CEAEnabled = "TEAMSFX_CEA_ENABLED";
+  static readonly MultiTenant = "TEAMSFX_MULTI_TENANT";
 }
 
 export interface FeatureFlag {
@@ -109,6 +110,10 @@ export class FeatureFlags {
     name: FeatureFlagName.CEAEnabled,
     defaultValue: "false",
   };
+  static readonly MultiTenant = {
+    name: FeatureFlagName.MultiTenant,
+    defaultValue: "false",
+  };
 }
 
 export function isCopilotExtensionEnabled(): boolean {

packages/fx-core/src/common/tools.ts

@@ -14,6 +14,22 @@ export async function getSideloadingStatus(token: string): Promise<boolean | und
   return teamsDevPortalClient.getSideloadingStatus(token);
 }
 
+export async function listAllTenants(token: string): Promise<Record<string, any>[]> {
+  const RM_ENDPOINT = "https://management.azure.com/tenants?api-version=2022-06-01";
+  if (token.length > 0) {
+    try {
+      const response = await axios.get(RM_ENDPOINT, {
+        headers: { Authorization: `Bearer ${token}` },
+      });
+      return response.data.value;
+    } catch (error) {
+      return [];
+    }
+  }
+
+  return [];
+}
+
 export async function getSPFxTenant(graphToken: string): Promise<string> {
   const GRAPH_TENANT_ENDPT = "https://graph.microsoft.com/v1.0/sites/root?$select=webUrl";
   if (graphToken.length > 0) {

packages/fx-core/tests/common/tools.test.ts

@@ -12,7 +12,12 @@ import * as path from "path";
 import Sinon, * as sinon from "sinon";
 import { getProjectMetadata } from "../../src/common/projectSettingsHelper";
 import * as telemetry from "../../src/common/telemetry";
-import { getSPFxToken, getSideloadingStatus, listDevTunnels } from "../../src/common/tools";
+import {
+  getSPFxToken,
+  getSideloadingStatus,
+  listAllTenants,
+  listDevTunnels,
+} from "../../src/common/tools";
 import { PackageService } from "../../src/component/m365/packageService";
 import { isVideoFilterProject } from "../../src/core/middleware/videoFilterAppBlocker";
 import { isUserCancelError } from "../../src/error/common";
@@ -124,6 +129,52 @@ describe("tools", () => {
     });
   });
 
+  describe("listAllTenants", () => {
+    const sandbox = sinon.createSandbox();
+
+    afterEach(() => {
+      sandbox.restore();
+    });
+
+    it("returns empty for invalid token", async () => {
+      const tenants = await listAllTenants("");
+
+      chai.assert.equal(tenants.length, 0);
+    });
+
+    it("returns empty when API call failure", async () => {
+      sandbox.stub(axios, "get").throws({ name: 404, message: "failed" });
+
+      const tenants = await listAllTenants("faked token");
+
+      chai.assert.equal(tenants.length, 0);
+    });
+
+    it("returns tenant list", async () => {
+      const fakedTenants = {
+        data: {
+          value: [
+            {
+              tenantId: "0022fd51-06f5-4557-8a34-69be98de6e20",
+              countryCode: "SG",
+              displayName: "MSFT",
+            },
+            {
+              tenantId: "313ef12c-d7cb-4f01-af90-1b113db5aa9a",
+              countryCode: "CN",
+              displayName: "Cisco",
+            },
+          ],
+        },
+      };
+      sandbox.stub(axios, "get").resolves(fakedTenants);
+
+      const tenants = await listAllTenants("faked token");
+
+      chai.assert.equal(tenants, fakedTenants.data.value);
+    });
+  });
+
   describe("getCopilotStatus", () => {
     let mockGet: () => AxiosResponse;
     let errors: number;

packages/vscode-extension/src/handlers/accounts/signinAccountHandlers.ts

@@ -34,7 +34,7 @@ export async function signinM365Callback(...args: unknown[]): Promise<Result<nul
   });
   const token = tokenRes.isOk() ? tokenRes.value : undefined;
   if (token !== undefined && node) {
-    node.setSignedIn((token as any).upn ? (token as any).upn : "");
+    await node.setSignedIn((token as any).upn ? (token as any).upn : "", (token as any).tid ?? "");
   }
 
   await envTreeProviderInstance.reloadEnvironments();

packages/vscode-extension/src/treeview/account/accountTreeViewProvider.ts

@@ -72,8 +72,9 @@ async function m365AccountStatusChangeHandler(
   const instance = AccountTreeViewProvider.getInstance();
   if (status === "SignedIn") {
     if (accountInfo) {
-      instance.m365AccountNode.setSignedIn(
-        (accountInfo.upn as string) ? (accountInfo.upn as string) : ""
+      await instance.m365AccountNode.setSignedIn(
+        (accountInfo.upn as string) ? (accountInfo.upn as string) : "",
+        (accountInfo.tid as string) ?? ""
       );
       if (token && source === "appStudio") {
         instance.m365AccountNode.updateChecks(token, true, true);
@@ -99,7 +100,11 @@ async function azureAccountStatusChangeHandler(
   if (status === "SignedIn") {
     const username = (accountInfo?.email as string) || (accountInfo?.upn as string);
     if (username) {
-      instance.azureAccountNode.setSignedIn(username);
+      await instance.azureAccountNode.setSignedIn(
+        token as string,
+        accountInfo?.tid as string,
+        username
+      );
       await envTreeProviderInstance.reloadEnvironments();
     }
   } else if (status === "SigningIn") {

packages/vscode-extension/src/treeview/account/azureNode.ts

@@ -7,6 +7,8 @@ import { TelemetryTriggerFrom } from "../../telemetry/extTelemetryEvents";
 import { localize } from "../../utils/localizeUtils";
 import { DynamicNode } from "../dynamicNode";
 import { AccountItemStatus, azureIcon, loadingIcon } from "./common";
+import { featureFlagManager, FeatureFlags } from "@microsoft/teamsfx-core";
+import { listAllTenants } from "@microsoft/teamsfx-core/build/common/tools";
 
 export class AzureAccountNode extends DynamicNode {
   public status: AccountItemStatus;
@@ -17,12 +19,20 @@ export class AzureAccountNode extends DynamicNode {
     this.contextValue = "signinAzure";
   }
 
-  public setSignedIn(upn: string) {
+  public async setSignedIn(token: string, tid: string, upn: string) {
     if (this.status === AccountItemStatus.SignedIn && this.label === upn) {
       return false;
     }
     this.status = AccountItemStatus.SignedIn;
     this.label = upn;
+    if (featureFlagManager.getBooleanValue(FeatureFlags.MultiTenant)) {
+      const tenants = await listAllTenants(token);
+      for (const tenant of tenants) {
+        if (tenant.tenantId === tid && tenant.displayName) {
+          this.label = `${upn} (${tenant.displayName as string})`;
+        }
+      }
+    }
     this.contextValue = "signedinAzure";
     this.eventEmitter.fire(undefined);
     return false;

packages/vscode-extension/src/treeview/account/m365Node.ts

@@ -1,14 +1,16 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { featureFlagManager, FeatureFlags as FxCoreFeatureFlags } from "@microsoft/teamsfx-core";
+import { AzureScopes, featureFlagManager, FeatureFlags } from "@microsoft/teamsfx-core";
 import * as vscode from "vscode";
 import { TelemetryTriggerFrom } from "../../telemetry/extTelemetryEvents";
 import { localize } from "../../utils/localizeUtils";
 import { DynamicNode } from "../dynamicNode";
 import { AccountItemStatus, loadingIcon, m365Icon } from "./common";
 import { CopilotNode } from "./copilotNode";
 import { SideloadingNode } from "./sideloadingNode";
+import { tools } from "../../globalVariables";
+import { listAllTenants } from "@microsoft/teamsfx-core/build/common/tools";
 
 export class M365AccountNode extends DynamicNode {
   public status: AccountItemStatus;
@@ -23,12 +25,26 @@ export class M365AccountNode extends DynamicNode {
     this.copilotNode = new CopilotNode(this.eventEmitter, "");
   }
 
-  public setSignedIn(upn: string) {
+  public async setSignedIn(upn: string, tid: string) {
     if (this.status === AccountItemStatus.SignedIn) {
       return;
     }
     this.status = AccountItemStatus.SignedIn;
+
     this.label = upn;
+    if (featureFlagManager.getBooleanValue(FeatureFlags.MultiTenant)) {
+      const tokenRes = await tools.tokenProvider.m365TokenProvider.getAccessToken({
+        scopes: AzureScopes,
+      });
+      if (tokenRes.isOk() && tokenRes.value) {
+        const tenants = await listAllTenants(tokenRes.value);
+        for (const tenant of tenants) {
+          if (tenant.tenantId === tid && tenant.displayName) {
+            this.label = `${upn} (${tenant.displayName as string})`;
+          }
+        }
+      }
+    }
     this.contextValue = "signedinM365";
     // refresh
     this.eventEmitter.fire(undefined);

packages/vscode-extension/test/handlers/accounts/signinAccountHandlers.test.ts

@@ -1,19 +1,15 @@
 import * as sinon from "sinon";
 import * as chai from "chai";
 import * as vscode from "vscode";
-import { UserCancelError } from "@microsoft/teamsfx-core";
+import { NetworkError, UserCancelError } from "@microsoft/teamsfx-core";
 import { AzureAccountManager } from "../../../src/commonlib/azureLogin";
 import {
   signinAzureCallback,
   signinM365Callback,
 } from "../../../src/handlers/accounts/signinAccountHandlers";
 import { ExtTelemetry } from "../../../src/telemetry/extTelemetry";
 import { setTools, tools } from "../../../src/globalVariables";
-import { ok } from "@microsoft/teamsfx-api";
-import VsCodeLogInstance from "../../../src/commonlib/log";
-import { VsCodeUI } from "../../../src/qm/vsc_ui";
-import { getExpService } from "../../../src/exp";
-import M365TokenInstance from "../../../src/commonlib/m365Login";
+import { err, ok } from "@microsoft/teamsfx-api";
 import { MockTools } from "../../mocks/mockTools";
 
 describe("SigninAccountHandlers", () => {
@@ -85,12 +81,72 @@ describe("SigninAccountHandlers", () => {
       sandbox.stub(ExtTelemetry, "sendTelemetryEvent");
     });
 
-    it("Happy path", async () => {
+    it("Happy path - valid upn", async () => {
       const setSignedInStub = sandbox.stub();
       const getJsonObjectStub = sandbox
         .stub(tools.tokenProvider.m365TokenProvider, "getJsonObject")
         .returns(Promise.resolve(ok({ upn: "test" })));
 
+      await signinM365Callback(
+        {},
+        {
+          status: 0,
+          setSignedIn: (...args: any[]) => {
+            setSignedInStub(args);
+          },
+        }
+      );
+
+      chai.assert.isTrue(getJsonObjectStub.calledOnce);
+      chai.assert.isTrue(setSignedInStub.calledOnceWith(["test", ""]));
+    });
+
+    it("Happy path - valid tid", async () => {
+      const setSignedInStub = sandbox.stub();
+      const getJsonObjectStub = sandbox
+        .stub(tools.tokenProvider.m365TokenProvider, "getJsonObject")
+        .returns(Promise.resolve(ok({ tid: "test" })));
+
+      await signinM365Callback(
+        {},
+        {
+          status: 0,
+          setSignedIn: (...args: any[]) => {
+            setSignedInStub(args);
+          },
+        }
+      );
+
+      chai.assert.isTrue(getJsonObjectStub.calledOnce);
+      chai.assert.isTrue(setSignedInStub.calledOnceWith(["", "test"]));
+    });
+
+    it("Happy path - valid upn & tid", async () => {
+      const setSignedInStub = sandbox.stub();
+      const getJsonObjectStub = sandbox
+        .stub(tools.tokenProvider.m365TokenProvider, "getJsonObject")
+        .returns(Promise.resolve(ok({ upn: "test upn", tid: "test tid" })));
+
+      await signinM365Callback(
+        {},
+        {
+          status: 0,
+          setSignedIn: (...args: any[]) => {
+            setSignedInStub(args);
+          },
+        }
+      );
+
+      chai.assert.isTrue(getJsonObjectStub.calledOnce);
+      chai.assert.isTrue(setSignedInStub.calledOnceWith(["test upn", "test tid"]));
+    });
+
+    it("invalid token result", async () => {
+      const setSignedInStub = sandbox.stub();
+      const getJsonObjectStub = sandbox
+        .stub(tools.tokenProvider.m365TokenProvider, "getJsonObject")
+        .returns(Promise.resolve(err(new NetworkError("source", "Failed to retrieve token"))));
+
       await signinM365Callback(
         {},
         {
@@ -102,7 +158,7 @@ describe("SigninAccountHandlers", () => {
       );
 
       chai.assert.isTrue(getJsonObjectStub.calledOnce);
-      chai.assert.isTrue(setSignedInStub.calledOnceWith("test"));
+      chai.assert.isTrue(setSignedInStub.notCalled);
     });
 
     it("Signed in", async () => {

packages/vscode-extension/test/treeview/account/accountsTreeViewProvider.test.ts

@@ -80,7 +80,19 @@ describe("AccountTreeViewProvider", () => {
     const m365SignedInStub = sandbox.stub(AccountTreeViewProvider.m365AccountNode, "setSignedIn");
     const updateChecksStub = sandbox.stub(AccountTreeViewProvider.m365AccountNode, "updateChecks");
     await m365StatusChange("SignedIn", "token", { upn: "upn" });
-    chai.assert.isTrue(m365SignedInStub.calledOnce);
+    chai.assert.isTrue(m365SignedInStub.calledOnceWithExactly("upn", ""));
+    chai.assert.isTrue(updateChecksStub.calledOnce);
+
+    m365SignedInStub.reset();
+    updateChecksStub.reset();
+    await m365StatusChange("SignedIn", "token", { tid: "tid" });
+    chai.assert.isTrue(m365SignedInStub.calledOnceWithExactly("", "tid"));
+    chai.assert.isTrue(updateChecksStub.calledOnce);
+
+    m365SignedInStub.reset();
+    updateChecksStub.reset();
+    await m365StatusChange("SignedIn", "token", { upn: "upn", tid: "tid" });
+    chai.assert.isTrue(m365SignedInStub.calledOnceWithExactly("upn", "tid"));
     chai.assert.isTrue(updateChecksStub.calledOnce);
 
     const m365SwitchingStub = sandbox.stub(AccountTreeViewProvider.m365AccountNode, "setSwitching");
@@ -98,6 +110,10 @@ describe("AccountTreeViewProvider", () => {
     await azureStatusChange("SignedIn", "token", { upn: "upn" });
     chai.assert.isTrue(azureSignedInStub.calledOnce);
 
+    azureSignedInStub.reset();
+    await azureStatusChange("SignedIn", "token", { upn: "upn", tid: "tid" });
+    chai.assert.isTrue(azureSignedInStub.calledOnceWithExactly("token", "tid", "upn"));
+
     const azureSigningInStub = sandbox.stub(
       AccountTreeViewProvider.azureAccountNode,
       "setSigningIn"