Add ability to present Managed Identity derived bearer token

src/environment.py

@@ -0,0 +1,35 @@
+import os
+from enum import Enum
+
+
+class Envs(Enum):
+    DEVELOPMENT = "dev"
+    REVIEW = "review"
+    PREPROD = "preprod"
+    PRODUCTION = "prod"
+
+
+class Environment:
+    @property
+    def development(self) -> bool:
+        return self.environment == Envs.DEVELOPMENT.value
+
+    @property
+    def production(self) -> bool:
+        return self.environment == Envs.PRODUCTION.value
+
+    @property
+    def review(self) -> bool:
+        return self.environment == Envs.REVIEW.value
+
+    @property
+    def preprod(self) -> bool:
+        return self.environment == Envs.PREPROD.value
+
+    @property
+    def environment(self) -> str:
+        env = os.getenv("ENVIRONMENT")
+        if not env or env.lower() not in (e.value for e in Envs):
+            return Envs.DEVELOPMENT.value
+        else:
+            return os.getenv("ENVIRONMENT", Envs.DEVELOPMENT.value).lower()

src/services/dicom/dicom_uploader.py

@@ -10,6 +10,9 @@
 from typing import Optional
 
 import requests
+from azure.identity import ManagedIdentityCredential
+
+from environment import Environment
 
 logger = logging.getLogger(__name__)
 
@@ -20,11 +23,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 +40,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 +58,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 | None:
+        resource = os.getenv("CLOUD_API_RESOURCE", "")
+        if resource or Environment().production:
+            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,43 @@ 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_in_production_with_no_cloud_api_resource(self, _, monkeypatch):
+        """Test that headers include access token from ManagedIdentityCredential in production even if CLOUD_API_RESOURCE is not set."""
+        monkeypatch.setenv("ENVIRONMENT", "prod")
+        with patch("services.dicom.dicom_uploader.ManagedIdentityCredential") as mock_credential:
+            mock_credential_instance = Mock()
+            mock_credential_instance.get_token.return_value.token = "prod_access_token"
+            mock_credential.return_value = mock_credential_instance
+
+            assert DICOMUploader().headers == {"Authorization": "Bearer prod_access_token"}
+            assert mock_credential_instance.get_token.call_args[0][0] == ""
+
+    def test_upload_headers_without_any_token(self, _, monkeypatch):
+        """Test that headers include empty token if neither CLOUD_API_RESOURCE nor CLOUD_API_TOKEN is set."""
+        monkeypatch.delenv("CLOUD_API_RESOURCE", raising=False)
+        monkeypatch.delenv("CLOUD_API_TOKEN", raising=False)
+        assert DICOMUploader().headers == {"Authorization": "Bearer "}

tests/test_environment.py

@@ -0,0 +1,40 @@
+from environment import Environment
+
+
+class TestEnvironment:
+    def test_environment(self, monkeypatch):
+
+        env = Environment()
+
+        # Default should be development
+        assert env.development
+        assert not env.production
+        assert not env.review
+        assert not env.preprod
+
+        # Test production environment
+        monkeypatch.setenv("ENVIRONMENT", "prod")
+        assert env.production
+        assert not env.development
+        assert not env.review
+        assert not env.preprod
+
+        # Test review environment
+        monkeypatch.setenv("ENVIRONMENT", "review")
+        assert env.review
+        assert not env.development
+        assert not env.production
+        assert not env.preprod
+
+        # Test preprod environment
+        monkeypatch.setenv("ENVIRONMENT", "preprod")
+        assert env.preprod
+        assert not env.development
+        assert not env.production
+
+        # Test unknown environment defaults to development
+        monkeypatch.setenv("ENVIRONMENT", "unknown")
+        assert env.development
+        assert not env.production
+        assert not env.review
+        assert not env.preprod