Tests

Author: heyitsaamir

packages/apps/src/microsoft/teams/apps/token_manager.py

@@ -130,12 +130,9 @@ async def _get_token_with_federated_identity(
         tenant_id_param = tenant_id or credentials.tenant_id or "botframework.com"
 
         # Step 1: Get MI token from api://AzureADTokenExchange
-        self._logger.debug("FIC Step 1: Acquiring MI token from api://AzureADTokenExchange")
         mi_token = await self._acquire_managed_identity_token(credentials)
-        self._logger.debug("FIC Step 1: Successfully acquired MI token")
 
         # Step 2: Use MI token as client_assertion to get final access token
-        self._logger.debug("FIC Step 2: Using MI token to acquire final access token")
         confidential_client = ConfidentialClientApplication(
             credentials.client_id,
             client_credential={"client_assertion": mi_token},
@@ -146,7 +143,6 @@ async def _get_token_with_federated_identity(
             lambda: confidential_client.acquire_token_for_client([scope])
         )
 
-        self._logger.debug("FIC Step 2: Successfully acquired final access token")
         return self._handle_token_response(token_res, error_prefix="FIC Step 2 failed")
 
     async def _acquire_managed_identity_token(self, credentials: FederatedIdentityCredentials) -> str:

packages/apps/tests/test_app.py

@@ -12,6 +12,7 @@
 from microsoft.teams.api import (
     Account,
     ConversationAccount,
+    FederatedIdentityCredentials,
     InvokeActivity,
     ManagedIdentityCredentials,
     MessageActivity,
@@ -639,23 +640,46 @@ def test_app_init_with_managed_identity(
             assert app.credentials.client_id == expected_client_id, f"Failed for: {description}"
             assert app.credentials.tenant_id == expected_tenant_id, f"Failed for: {description}"
 
-    def test_app_init_with_managed_identity_client_id_mismatch(self, mock_logger, mock_storage):
-        """Test app init raises error when managed_identity_client_id != client_id (federated identity)."""
-        # When managed_identity_client_id differs from client_id, should raise error
-        # (Federated Identity Credentials not yet supported)
+    @pytest.mark.parametrize(
+        "managed_identity_client_id,expected_mi_type,expected_mi_client_id,description",
+        [
+            # System-assigned managed identity
+            ("system", "system", None, "system-assigned managed identity"),
+            # User-assigned managed identity (federated)
+            (
+                "different-managed-identity-id",
+                "user",
+                "different-managed-identity-id",
+                "user-assigned federated identity",
+            ),
+        ],
+    )
+    def test_app_init_with_federated_identity(
+        self,
+        mock_logger,
+        mock_storage,
+        managed_identity_client_id: str,
+        expected_mi_type: str,
+        expected_mi_client_id: str | None,
+        description: str,
+    ):
+        """Test app initialization with FederatedIdentityCredentials."""
         options = AppOptions(
             logger=mock_logger,
             storage=mock_storage,
             client_id="app-client-id",
-            managed_identity_client_id="different-managed-identity-id",  # Different!
+            managed_identity_client_id=managed_identity_client_id,
         )
 
         with patch.dict("os.environ", {"CLIENT_SECRET": "", "TENANT_ID": "test-tenant-id"}, clear=False):
-            with pytest.raises(ValueError) as exc_info:
-                App(**options)
+            app = App(**options)
 
-            assert "Federated Identity Credentials is not yet supported" in str(exc_info.value)
-            assert "managed_identity_client_id must equal client_id" in str(exc_info.value)
+            assert app.credentials is not None, f"Failed for: {description}"
+            assert isinstance(app.credentials, FederatedIdentityCredentials), f"Failed for: {description}"
+            assert app.credentials.client_id == "app-client-id", f"Failed for: {description}"
+            assert app.credentials.managed_identity_type == expected_mi_type, f"Failed for: {description}"
+            assert app.credentials.managed_identity_client_id == expected_mi_client_id, f"Failed for: {description}"
+            assert app.credentials.tenant_id == "test-tenant-id", f"Failed for: {description}"
 
     def test_app_init_with_client_secret_takes_precedence(self, mock_logger, mock_storage):
         """Test that ClientCredentials is used when both client_secret and managed_identity_client_id are provided."""

packages/apps/tests/test_token_manager.py

@@ -3,10 +3,16 @@
 Licensed under the MIT License.
 """
 
+from typing import Literal, cast
 from unittest.mock import MagicMock, create_autospec, patch
 
 import pytest
-from microsoft.teams.api import ClientCredentials, JsonWebToken, ManagedIdentityCredentials
+from microsoft.teams.api import (
+    ClientCredentials,
+    FederatedIdentityCredentials,
+    JsonWebToken,
+    ManagedIdentityCredentials,
+)
 from microsoft.teams.apps.token_manager import TokenManager
 from msal import ManagedIdentityClient  # pyright: ignore[reportMissingTypeStubs]
 
@@ -128,8 +134,8 @@ async def test_get_token_with_managed_identity(self, get_token_method: str, expe
 
         manager = TokenManager(credentials=mock_credentials)
 
-        # Patch _get_msal_client to return our mock
-        with patch.object(manager, "_get_msal_client", return_value=mock_msal_client):
+        # Patch _get_managed_identity_client to return our mock
+        with patch.object(manager, "_get_managed_identity_client", return_value=mock_msal_client):
             # Call the method dynamically
             token = await getattr(manager, get_token_method)()
 
@@ -155,22 +161,22 @@ async def test_get_graph_token_with_managed_identity_and_tenant(self):
 
         manager = TokenManager(credentials=mock_credentials)
 
-        # Track calls to _get_msal_client
+        # Track calls to _get_managed_identity_client
         get_msal_client_calls: list[str] = []
 
-        def track_get_msal_client(tenant_id: str):
+        def track_get_managed_identity_client(credentials: ManagedIdentityCredentials, tenant_id: str):
             get_msal_client_calls.append(tenant_id)
             return mock_msal_client
 
-        # Patch _get_msal_client to track calls
-        with patch.object(manager, "_get_msal_client", side_effect=track_get_msal_client):
+        # Patch _get_managed_identity_client to track calls
+        with patch.object(manager, "_get_managed_identity_client", side_effect=track_get_managed_identity_client):
             # Request token for different tenant
             token = await manager.get_graph_token("different-tenant-id")
 
             assert token is not None
             assert isinstance(token, JsonWebToken)
 
-            # Verify _get_msal_client was called with different-tenant-id
+            # Verify _get_managed_identity_client was called with different-tenant-id
             assert "different-tenant-id" in get_msal_client_calls
 
     @pytest.mark.asyncio
@@ -190,10 +196,99 @@ async def test_get_token_error_handling_with_managed_identity(self):
 
         manager = TokenManager(credentials=mock_credentials)
 
-        # Patch _get_msal_client to return our mock
-        with patch.object(manager, "_get_msal_client", return_value=mock_msal_client):
+        # Patch _get_managed_identity_client to return our mock
+        with patch.object(manager, "_get_managed_identity_client", return_value=mock_msal_client):
             # Should raise an error when token acquisition fails
             with pytest.raises(ValueError) as exc_info:
                 await manager.get_bot_token()
 
             assert "invalid_client" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "mi_type,mi_client_id,description",
+        [
+            ("system", None, "system-assigned managed identity"),
+            ("user", "test-user-mi-client-id", "user-assigned managed identity"),
+        ],
+    )
+    async def test_get_token_with_federated_identity(self, mi_type: str, mi_client_id: str | None, description: str):
+        """Test token retrieval using FederatedIdentityCredentials (two-step flow)."""
+        mock_credentials = FederatedIdentityCredentials(
+            client_id="test-app-client-id",
+            managed_identity_type=cast(Literal["system", "user"], mi_type),
+            managed_identity_client_id=mi_client_id,
+            tenant_id="test-tenant-id",
+        )
+
+        manager = TokenManager(credentials=mock_credentials)
+
+        # Mock the managed identity token acquisition (step 1)
+        mi_token = "mi_token_from_step_1"
+        with patch.object(manager, "_acquire_managed_identity_token", return_value=mi_token):
+            # Mock ConfidentialClientApplication for step 2
+            with patch("microsoft.teams.apps.token_manager.ConfidentialClientApplication") as mock_confidential_app:
+                mock_app_instance = MagicMock()
+                mock_app_instance.acquire_token_for_client.return_value = {"access_token": VALID_TEST_TOKEN}
+                mock_confidential_app.return_value = mock_app_instance
+
+                token = await manager.get_bot_token()
+
+                assert token is not None, f"Failed for: {description}"
+                assert isinstance(token, JsonWebToken), f"Failed for: {description}"
+                assert str(token) == VALID_TEST_TOKEN, f"Failed for: {description}"
+
+                # Verify ConfidentialClientApplication was called with MI token as client_assertion
+                mock_confidential_app.assert_called_once()
+                call_kwargs = mock_confidential_app.call_args[1]
+                assert call_kwargs["client_credential"] == {"client_assertion": mi_token}, f"Failed for: {description}"
+
+    @pytest.mark.asyncio
+    async def test_get_token_with_federated_identity_step1_failure(self):
+        """Test error handling when step 1 (MI token acquisition) fails."""
+        mock_credentials = FederatedIdentityCredentials(
+            client_id="test-app-client-id",
+            managed_identity_type="user",
+            managed_identity_client_id="test-mi-client-id",
+            tenant_id="test-tenant-id",
+        )
+
+        manager = TokenManager(credentials=mock_credentials)
+
+        # Mock step 1 to fail
+        with patch.object(
+            manager, "_acquire_managed_identity_token", side_effect=ValueError("MI token acquisition failed")
+        ):
+            with pytest.raises(ValueError) as exc_info:
+                await manager.get_bot_token()
+
+            assert "MI token acquisition failed" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_get_token_with_federated_identity_step2_failure(self):
+        """Test error handling when step 2 (final token acquisition) fails."""
+        mock_credentials = FederatedIdentityCredentials(
+            client_id="test-app-client-id",
+            managed_identity_type="user",
+            managed_identity_client_id="test-mi-client-id",
+            tenant_id="test-tenant-id",
+        )
+
+        manager = TokenManager(credentials=mock_credentials)
+
+        # Mock step 1 to succeed
+        mi_token = "mi_token_from_step_1"
+        with patch.object(manager, "_acquire_managed_identity_token", return_value=mi_token):
+            # Mock step 2 to fail
+            with patch("microsoft.teams.apps.token_manager.ConfidentialClientApplication") as mock_confidential_app:
+                mock_app_instance = MagicMock()
+                mock_app_instance.acquire_token_for_client.return_value = {
+                    "error": "invalid_grant",
+                    "error_description": "FIC Step 2 failed",
+                }
+                mock_confidential_app.return_value = mock_app_instance
+
+                with pytest.raises(ValueError) as exc_info:
+                    await manager.get_bot_token()
+
+                assert "invalid_grant" in str(exc_info.value)