Add ability to present Managed Identity derived bearer token

Author: steventux

src/services/dicom/dicom_uploader.py

@@ -9,7 +9,9 @@
 import os
 from typing import Optional
 
+
 import requests
+from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
 
 logger = logging.getLogger(__name__)
 
@@ -20,11 +22,6 @@ def __init__(self, api_endpoint: str | None = None, timeout: int = 30, verify_ss
         self.timeout = timeout
         self.verify_ssl = verify_ssl
 
-    def headers(self) -> dict:
-        return {
-            "Authorization": f"Bearer {os.getenv('CLOUD_API_TOKEN', '')}",
-        }
-
     def upload_dicom(self, sop_instance_uid: str, dicom_stream: io.BufferedReader, action_id: Optional[str]) -> bool:
         if not action_id:
             logger.error(f"No action_id for {sop_instance_uid}, upload will be rejected by server")
@@ -42,7 +39,7 @@ def upload_dicom(self, sop_instance_uid: str, dicom_stream: io.BufferedReader, a
                 files=files,
                 timeout=self.timeout,
                 verify=self.verify_ssl,
-                headers=self.headers(),
+                headers=self.headers,
             )
 
             if response.status_code == 201:
@@ -60,3 +57,17 @@ def upload_dicom(self, sop_instance_uid: str, dicom_stream: io.BufferedReader, a
         except requests.exceptions.RequestException as e:
             logger.error(f"Upload error for {sop_instance_uid}: {e}", exc_info=True)
             return False
+
+    @property
+    def headers(self) -> dict:
+        return {
+            "Authorization": f"Bearer {self.access_token}",
+        }
+
+    @property
+    def access_token(self) -> str:
+        resource = os.getenv("CLOUD_API_RESOURCE")
+        if resource:
+            return ManagedIdentityCredential().get_token(resource).token
+        else:
+            return os.getenv("CLOUD_API_TOKEN", "")

tests/services/dicom/test_dicom_uploader.py

@@ -9,6 +9,7 @@
 from services.dicom.dicom_uploader import DICOMUploader
 
 
+@patch("services.dicom.dicom_uploader.requests.put")
 class TestDICOMUploader:
     @pytest.fixture
     def dicom_file(self):
@@ -17,7 +18,6 @@ def dicom_file(self):
         tf.close()
         yield tf.name
 
-    @patch("services.dicom.dicom_uploader.requests.put")
     def test_upload_success(self, mock_put, dicom_file):
         mock_response = Mock()
         mock_response.status_code = 201
@@ -39,7 +39,7 @@ def test_upload_success(self, mock_put, dicom_file):
             files=mock_put.call_args[1]["files"],
             timeout=30,
             verify=True,
-            headers=uploader.headers(),
+            headers=uploader.headers,
         )
 
         call_kwargs = mock_put.call_args[1]
@@ -49,14 +49,13 @@ def test_upload_success(self, mock_put, dicom_file):
         assert isinstance(file_tuple[1], io.BufferedReader)
         assert file_tuple[1].read() == open(dicom_file, "rb").read()
 
-    def test_upload_without_action_id(self, dicom_file):
+    def test_upload_without_action_id(self, _, dicom_file):
         """Upload without action_id does not make request."""
         uploader = DICOMUploader()
         result = uploader.upload_dicom(sop_instance_uid="1.2.3", dicom_stream=open(dicom_file, "rb"), action_id=None)
 
         assert result is False
 
-    @patch("services.dicom.dicom_uploader.requests.put")
     def test_upload_failure_status_code(self, mock_put, dicom_file):
         mock_response = Mock()
         mock_response.status_code = 500
@@ -68,7 +67,6 @@ def test_upload_failure_status_code(self, mock_put, dicom_file):
 
         assert result is False
 
-    @patch("services.dicom.dicom_uploader.requests.put")
     def test_upload_timeout(self, mock_put, dicom_file):
         mock_put.side_effect = requests.exceptions.Timeout()
 
@@ -77,11 +75,30 @@ def test_upload_timeout(self, mock_put, dicom_file):
 
         assert result is False
 
-    @patch("services.dicom.dicom_uploader.requests.put")
     def test_upload_network_error(self, mock_put, dicom_file):
         mock_put.side_effect = requests.exceptions.ConnectionError()
 
         uploader = DICOMUploader()
         result = uploader.upload_dicom("1.2.3", open(dicom_file, "rb"), None)
 
         assert result is False
+
+    def test_upload_headers_with_managed_identity_access_token(self, _, monkeypatch):
+        """Test that headers include access token from ManagedIdentityCredential."""
+        monkeypatch.setenv("CLOUD_API_RESOURCE", "https://example.com/.default")
+        with patch("services.dicom.dicom_uploader.ManagedIdentityCredential") as mock_credential:
+            mock_credential_instance = Mock()
+            mock_credential_instance.get_token.return_value.token = "fake_access_token"
+            mock_credential.return_value = mock_credential_instance
+
+            assert DICOMUploader().headers == {"Authorization": "Bearer fake_access_token"}
+
+    def test_upload_headers_without_managed_identity_resource(self, _, monkeypatch):
+        """Test that headers include CLOUD_API_TOKEN if CLOUD_API_RESOURCE is not set."""
+        monkeypatch.setenv("CLOUD_API_TOKEN", "env_access_token")
+
+        assert DICOMUploader().headers == {"Authorization": "Bearer env_access_token"}
+
+    def test_upload_headers_without_any_token(self, _):
+        """Test that headers include empty token if neither CLOUD_API_RESOURCE nor CLOUD_API_TOKEN is set."""
+        assert DICOMUploader().headers == {"Authorization": "Bearer "}