Add support for Azure AAD based Auth

tests/providers/test_azure_openai_provider.py

@@ -1,13 +1,18 @@
 """Test Azure OpenAI provider (Responses API via OpenAI SDK)."""
 
-from unittest.mock import AsyncMock, MagicMock
+import sys
+import time
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from nanobot.providers.azure_openai_provider import AzureOpenAIProvider
+from nanobot.providers.azure_openai_provider import (
+    AzureOpenAIProvider,
+    _AzureTokenProvider,
+)
 from nanobot.providers.base import LLMResponse
 
-
 # ---------------------------------------------------------------------------
 # Init & validation
 # ---------------------------------------------------------------------------
@@ -25,6 +30,8 @@ def test_init_creates_sdk_client():
     assert provider.default_model == "gpt-4o-deployment"
     # SDK client base_url ends with /openai/v1/
     assert str(provider._client.base_url).rstrip("/").endswith("/openai/v1")
+    # Static-key path must NOT construct an AAD token provider
+    assert provider._token_provider is None
 
 
 def test_init_base_url_no_trailing_slash():
@@ -42,11 +49,6 @@ def test_init_base_url_with_trailing_slash():
     assert str(provider._client.base_url).rstrip("/").endswith("/openai/v1")
 
 
-def test_init_validation_missing_key():
-    with pytest.raises(ValueError, match="Azure OpenAI api_key is required"):
-        AzureOpenAIProvider(api_key="", api_base="https://test.com")
-
-
 def test_init_validation_missing_base():
     with pytest.raises(ValueError, match="Azure OpenAI api_base is required"):
         AzureOpenAIProvider(api_key="test", api_base="")
@@ -59,6 +61,96 @@ def test_no_api_version_in_base_url():
     assert "api-version" not in base
 
 
+# ---------------------------------------------------------------------------
+# AAD / DefaultAzureCredential fallback
+# ---------------------------------------------------------------------------
+
+
+def _install_fake_azure_identity(monkeypatch, credential_factory):
+    """Install a fake ``azure.identity.aio`` module exposing ``DefaultAzureCredential``."""
+    azure_mod = sys.modules.get("azure") or SimpleNamespace()
+    identity_mod = SimpleNamespace()
+    aio_mod = SimpleNamespace(DefaultAzureCredential=credential_factory)
+    identity_mod.aio = aio_mod
+    azure_mod.identity = identity_mod  # type: ignore[attr-defined]
+
+    monkeypatch.setitem(sys.modules, "azure", azure_mod)
+    monkeypatch.setitem(sys.modules, "azure.identity", identity_mod)
+    monkeypatch.setitem(sys.modules, "azure.identity.aio", aio_mod)
+
+
+def test_init_missing_key_uses_aad_token_provider(monkeypatch):
+    """Empty api_key falls back to DefaultAzureCredential via _AzureTokenProvider."""
+    credential_instance = MagicMock()
+    credential_factory = MagicMock(return_value=credential_instance)
+    _install_fake_azure_identity(monkeypatch, credential_factory)
+
+    provider = AzureOpenAIProvider(
+        api_key="", api_base="https://res.openai.azure.com",
+    )
+
+    assert provider._token_provider is not None
+    assert isinstance(provider._token_provider, _AzureTokenProvider)
+    # DefaultAzureCredential must have been instantiated exactly once
+    credential_factory.assert_called_once_with()
+    # The SDK client must have received the token provider as its api_key
+    # (the SDK stores it on the auth wrapper, not directly accessible — so
+    # we assert the callable was wired in via the provider attribute).
+    assert provider._token_provider._credential is credential_instance
+
+
+def test_init_explicit_key_does_not_construct_credential(monkeypatch):
+    """Explicit api_key wins; DefaultAzureCredential must never be touched."""
+    credential_factory = MagicMock(side_effect=AssertionError(
+        "DefaultAzureCredential must not be constructed when api_key is set"
+    ))
+    _install_fake_azure_identity(monkeypatch, credential_factory)
+
+    provider = AzureOpenAIProvider(
+        api_key="real-key", api_base="https://res.openai.azure.com",
+    )
+
+    assert provider._token_provider is None
+    credential_factory.assert_not_called()
+
+
+def test_init_missing_key_without_azure_identity_raises(monkeypatch):
+    """Clear RuntimeError with pip-install hint when azure-identity is missing."""
+    # Force the import inside _AzureTokenProvider to fail.
+    real_import = __builtins__["__import__"] if isinstance(__builtins__, dict) else __builtins__.__import__
+
+    def fake_import(name, *args, **kwargs):
+        if name == "azure.identity.aio":
+            raise ImportError("No module named 'azure.identity.aio'")
+        return real_import(name, *args, **kwargs)
+
+    with patch("builtins.__import__", side_effect=fake_import):
+        with pytest.raises(RuntimeError, match=r"pip install 'nanobot-ai\[azure\]'"):
+            AzureOpenAIProvider(api_key="", api_base="https://res.openai.azure.com")
+
+
+@pytest.mark.asyncio
+async def test_token_provider_returns_credential_token(monkeypatch):
+    """Token callback delegates to DefaultAzureCredential.get_token with the AOAI scope."""
+    access_token = SimpleNamespace(token="token-A", expires_on=time.time() + 3600)
+    credential_instance = MagicMock()
+    credential_instance.get_token = AsyncMock(return_value=access_token)
+
+    credential_factory = MagicMock(return_value=credential_instance)
+    _install_fake_azure_identity(monkeypatch, credential_factory)
+
+    tp = _AzureTokenProvider()
+
+    assert await tp() == "token-A"
+    credential_instance.get_token.assert_awaited_with(
+        "https://cognitiveservices.azure.com/.default"
+    )
+    # No client-side caching layer — every call delegates to the Azure SDK,
+    # which has its own MSAL-backed token cache.
+    assert await tp() == "token-A"
+    assert credential_instance.get_token.await_count == 2
+
+
 # ---------------------------------------------------------------------------
 # _supports_temperature
 # ---------------------------------------------------------------------------