Skip to content

Commit 86a1962

Browse files
Merge pull request #236 from microsoft/dev
feat: Use ManagedIdentityCredential fir Prod Workloads and DefaultAzureCredential for dev
2 parents 2643dcc + 06b2f12 commit 86a1962

File tree

26 files changed

+419
-40
lines changed

26 files changed

+419
-40
lines changed

infra/main.bicep

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -810,6 +810,10 @@ module avmContainerApp 'br/public:avm/res/app/container-app:0.17.0' = {
810810
name: 'APP_CONFIG_ENDPOINT'
811811
value: ''
812812
}
813+
{
814+
name: 'APP_ENV'
815+
value: 'prod'
816+
}
813817
]
814818
}
815819
]
@@ -854,6 +858,10 @@ module avmContainerApp_API 'br/public:avm/res/app/container-app:0.17.0' = {
854858
name: 'APP_CONFIG_ENDPOINT'
855859
value: ''
856860
}
861+
{
862+
name: 'APP_ENV'
863+
value: 'prod'
864+
}
857865
]
858866
probes: [
859867
// Liveness Probe - Checks if the app is still running
@@ -1272,6 +1280,10 @@ module avmContainerApp_update 'br/public:avm/res/app/container-app:0.17.0' = {
12721280
name: 'APP_CONFIG_ENDPOINT'
12731281
value: avmAppConfig.outputs.endpoint
12741282
}
1283+
{
1284+
name: 'APP_ENV'
1285+
value: 'prod'
1286+
}
12751287
]
12761288
}
12771289
]
@@ -1327,6 +1339,10 @@ module avmContainerApp_API_update 'br/public:avm/res/app/container-app:0.17.0' =
13271339
name: 'APP_CONFIG_ENDPOINT'
13281340
value: avmAppConfig.outputs.endpoint
13291341
}
1342+
{
1343+
name: 'APP_ENV'
1344+
value: 'prod'
1345+
}
13301346
]
13311347
probes: [
13321348
// Liveness Probe - Checks if the app is still running

src/ContentProcessor/conftest.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
"""
2+
Global test configuration and fixtures for ContentProcessor tests.
3+
"""
4+
import sys
5+
import os
6+
import pytest
7+
from unittest.mock import patch, MagicMock
8+
9+
# Add src directory to Python path for imports
10+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
11+
12+
pytest_plugins = ["pytest_mock"]
13+
14+
15+
@pytest.fixture(autouse=True, scope="function")
16+
def mock_azure_credentials_for_helpers(request):
17+
"""
18+
Mock Azure credentials for azure_helper classes only.
19+
Skip this for credential utility tests that need to test the actual logic.
20+
"""
21+
# Skip mocking for credential utility tests
22+
if "test_azure_credential_utils" in str(request.fspath):
23+
yield
24+
return
25+
26+
with patch("helpers.azure_credential_utils.get_azure_credential") as mock_get_cred, \
27+
patch("helpers.azure_credential_utils.get_azure_credential_async") as mock_get_cred_async:
28+
29+
# Create mock credential objects
30+
mock_credential = MagicMock()
31+
mock_get_cred.return_value = mock_credential
32+
mock_get_cred_async.return_value = mock_credential
33+
34+
yield mock_credential

src/ContentProcessor/pytest.ini

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
[tool:pytest]
2+
addopts = -v --strict-markers --disable-warnings
3+
python_files = tests.py test_*.py *_tests.py
4+
testpaths = src/tests
5+
markers =
6+
asyncio: marks tests as async
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT License.
3+
4+
import os
5+
from azure.identity import ManagedIdentityCredential, DefaultAzureCredential
6+
from azure.identity.aio import ManagedIdentityCredential as AioManagedIdentityCredential, DefaultAzureCredential as AioDefaultAzureCredential
7+
8+
9+
async def get_azure_credential_async(client_id=None):
10+
"""
11+
Returns an Azure credential asynchronously based on the application environment.
12+
13+
If the environment is 'dev', it uses AioDefaultAzureCredential.
14+
Otherwise, it uses AioManagedIdentityCredential.
15+
16+
Args:
17+
client_id (str, optional): The client ID for the Managed Identity Credential.
18+
19+
Returns:
20+
Credential object: Either AioDefaultAzureCredential or AioManagedIdentityCredential.
21+
"""
22+
if os.getenv("APP_ENV", "prod").lower() == 'dev':
23+
return AioDefaultAzureCredential() # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
24+
else:
25+
return AioManagedIdentityCredential(client_id=client_id)
26+
27+
28+
def get_azure_credential(client_id=None):
29+
"""
30+
Returns an Azure credential based on the application environment.
31+
32+
If the environment is 'dev', it uses DefaultAzureCredential.
33+
Otherwise, it uses ManagedIdentityCredential.
34+
35+
Args:
36+
client_id (str, optional): The client ID for the Managed Identity Credential.
37+
38+
Returns:
39+
Credential object: Either DefaultAzureCredential or ManagedIdentityCredential.
40+
"""
41+
if os.getenv("APP_ENV", "prod").lower() == 'dev':
42+
return DefaultAzureCredential() # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
43+
else:
44+
return ManagedIdentityCredential(client_id=client_id)

src/ContentProcessor/src/libs/application/application_context.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
from azure.identity import DefaultAzureCredential
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT License.
3+
from typing import Any
24

35
from libs.application.application_configuration import AppConfiguration
46
from libs.base.application_models import AppModelBase
@@ -11,10 +13,10 @@ class AppContext(AppModelBase):
1113
"""
1214

1315
configuration: AppConfiguration = None
14-
credential: DefaultAzureCredential = None
16+
credential: Any = None # Azure credential object
1517

1618
def set_configuration(self, configuration: AppConfiguration):
1719
self.configuration = configuration
1820

19-
def set_credential(self, credential: DefaultAzureCredential):
21+
def set_credential(self, credential: Any):
2022
self.credential = credential

src/ContentProcessor/src/libs/azure_helper/app_configuration.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,18 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT License.
3+
14
import os
25

36
from azure.appconfiguration import AzureAppConfigurationClient
4-
from azure.identity import DefaultAzureCredential
7+
from helpers.azure_credential_utils import get_azure_credential
58

69

710
class AppConfigurationHelper:
8-
credential: DefaultAzureCredential = None
911
app_config_endpoint: str = None
1012
app_config_client: AzureAppConfigurationClient = None
1113

1214
def __init__(self, app_config_endpoint: str):
13-
self.credential = DefaultAzureCredential()
15+
self.credential = get_azure_credential()
1416
self.app_config_endpoint = app_config_endpoint
1517
self._initialize_client()
1618

src/ContentProcessor/src/libs/azure_helper/azure_openai.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
1-
from azure.identity import DefaultAzureCredential, get_bearer_token_provider
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT License.
3+
4+
from azure.identity import get_bearer_token_provider
5+
from helpers.azure_credential_utils import get_azure_credential
26
from openai import AzureOpenAI
37

48

59
def get_openai_client(azure_openai_endpoint: str) -> AzureOpenAI:
6-
credential = DefaultAzureCredential()
10+
credential = get_azure_credential()
711
token_provider = get_bearer_token_provider(
812
credential, "https://cognitiveservices.azure.com/.default"
913
)

src/ContentProcessor/src/libs/azure_helper/content_understanding.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,22 +7,21 @@
77
from pathlib import Path
88

99
import requests
10-
from azure.identity import DefaultAzureCredential
10+
from helpers.azure_credential_utils import get_azure_credential
1111
from requests.models import Response
1212

1313
COGNITIVE_SERVICES_SCOPE = "https://cognitiveservices.azure.com/.default"
1414

1515

1616
class AzureContentUnderstandingHelper:
17-
credential: DefaultAzureCredential = None
1817

1918
def __init__(
2019
self,
2120
endpoint: str,
2221
api_version: str = "2024-12-01-preview",
2322
x_ms_useragent: str = "cps-contentunderstanding/client",
2423
):
25-
self.credential = DefaultAzureCredential()
24+
self.credential = get_azure_credential()
2625

2726
if not api_version:
2827
raise ValueError("API version must be provided.")

src/ContentProcessor/src/libs/azure_helper/storage_blob.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,19 @@
33

44
from typing import IO, Union
55

6-
from azure.identity import DefaultAzureCredential
6+
from helpers.azure_credential_utils import get_azure_credential
77
from azure.storage.blob import BlobServiceClient
88

99

1010
class StorageBlobHelper:
11-
credential: DefaultAzureCredential = None
1211
blob_service_client: BlobServiceClient = None
1312

1413
@staticmethod
1514
def get(account_url: str, container_name: str = None):
1615
return StorageBlobHelper(account_url=account_url, container_name=container_name)
1716

1817
def __init__(self, account_url: str, container_name=None):
19-
self.credential = DefaultAzureCredential()
18+
self.credential = get_azure_credential()
2019
self.blob_service_client = BlobServiceClient(
2120
account_url=account_url, credential=self.credential
2221
)

src/ContentProcessor/src/libs/pipeline/pipeline_queue_helper.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import logging
55

66
from azure.core.exceptions import ResourceNotFoundError
7-
from azure.identity import DefaultAzureCredential
7+
from helpers.azure_credential_utils import get_azure_credential
88
from azure.storage.queue import QueueClient, QueueMessage
99

1010
from libs.pipeline import pipeline_step_helper
@@ -28,7 +28,7 @@ def invalidate_queue(queue_client: QueueClient):
2828

2929

3030
def create_or_get_queue_client(
31-
queue_name: str, accouont_url: str, credential: DefaultAzureCredential
31+
queue_name: str, accouont_url: str, credential: get_azure_credential
3232
) -> QueueClient:
3333
queue_client = QueueClient(
3434
account_url=accouont_url, queue_name=queue_name, credential=credential
@@ -55,7 +55,7 @@ def has_messages(queue_client: QueueClient) -> bool:
5555

5656

5757
def pass_data_pipeline_to_next_step(
58-
data_pipeline: DataPipeline, account_url: str, credential: DefaultAzureCredential
58+
data_pipeline: DataPipeline, account_url: str, credential: get_azure_credential
5959
):
6060
next_step_name = pipeline_step_helper.get_next_step_name(
6161
data_pipeline.pipeline_status, data_pipeline.pipeline_status.active_step
@@ -70,7 +70,7 @@ def pass_data_pipeline_to_next_step(
7070

7171

7272
def _create_queue_client(
73-
account_url: str, queue_name: str, credential: DefaultAzureCredential
73+
account_url: str, queue_name: str, credential: get_azure_credential
7474
) -> QueueClient:
7575
queue_client = QueueClient(
7676
account_url=account_url, queue_name=queue_name, credential=credential

0 commit comments

Comments
 (0)