From 3e8c34a59805a83528715495a6bdc1cad025eaa6 Mon Sep 17 00:00:00 2001 From: Steve Laing Date: Wed, 15 Apr 2026 12:53:18 +0100 Subject: [PATCH 1/3] Add ability to present Managed Identity derived bearer token --- src/services/dicom/dicom_uploader.py | 23 +++++++++++----- tests/services/dicom/test_dicom_uploader.py | 29 ++++++++++++++++----- 2 files changed, 40 insertions(+), 12 deletions(-) diff --git a/src/services/dicom/dicom_uploader.py b/src/services/dicom/dicom_uploader.py index 2f5365b4..8c8cd1c3 100644 --- a/src/services/dicom/dicom_uploader.py +++ b/src/services/dicom/dicom_uploader.py @@ -9,7 +9,9 @@ import os from typing import Optional + import requests +from azure.identity import 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", "") diff --git a/tests/services/dicom/test_dicom_uploader.py b/tests/services/dicom/test_dicom_uploader.py index 9c6614e2..2d3d4d60 100644 --- a/tests/services/dicom/test_dicom_uploader.py +++ b/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,7 +75,6 @@ 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() @@ -85,3 +82,23 @@ def test_upload_network_error(self, mock_put, dicom_file): 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 "} From 970b8b4cad36aade4c980510563196412baa02b9 Mon Sep 17 00:00:00 2001 From: Steve Laing Date: Thu, 23 Apr 2026 16:39:39 +0100 Subject: [PATCH 2/3] Add Environment helper --- src/environment.py | 35 ++++++++++++++++++++++++++++++++++ tests/test_environment.py | 40 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+) create mode 100644 src/environment.py create mode 100644 tests/test_environment.py diff --git a/src/environment.py b/src/environment.py new file mode 100644 index 00000000..082870b6 --- /dev/null +++ b/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() diff --git a/tests/test_environment.py b/tests/test_environment.py new file mode 100644 index 00000000..1cc21968 --- /dev/null +++ b/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 From e1255a969182c1a551b704509d4d9800235b67b0 Mon Sep 17 00:00:00 2001 From: Steve Laing Date: Thu, 23 Apr 2026 16:40:40 +0100 Subject: [PATCH 3/3] Always present MI token in production environment --- src/services/dicom/dicom_uploader.py | 9 +++++---- tests/services/dicom/test_dicom_uploader.py | 15 ++++++++++++++- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/services/dicom/dicom_uploader.py b/src/services/dicom/dicom_uploader.py index 8c8cd1c3..f52754b2 100644 --- a/src/services/dicom/dicom_uploader.py +++ b/src/services/dicom/dicom_uploader.py @@ -9,10 +9,11 @@ import os from typing import Optional - import requests from azure.identity import ManagedIdentityCredential +from environment import Environment + logger = logging.getLogger(__name__) @@ -65,9 +66,9 @@ def headers(self) -> dict: } @property - def access_token(self) -> str: - resource = os.getenv("CLOUD_API_RESOURCE") - if resource: + 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", "") diff --git a/tests/services/dicom/test_dicom_uploader.py b/tests/services/dicom/test_dicom_uploader.py index 2d3d4d60..65ce5047 100644 --- a/tests/services/dicom/test_dicom_uploader.py +++ b/tests/services/dicom/test_dicom_uploader.py @@ -99,6 +99,19 @@ def test_upload_headers_without_managed_identity_resource(self, _, monkeypatch): assert DICOMUploader().headers == {"Authorization": "Bearer env_access_token"} - def test_upload_headers_without_any_token(self, _): + 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 "}