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/src/services/dicom/dicom_uploader.py b/src/services/dicom/dicom_uploader.py index 2f5365b4..f52754b2 100644 --- a/src/services/dicom/dicom_uploader.py +++ b/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", "") diff --git a/tests/services/dicom/test_dicom_uploader.py b/tests/services/dicom/test_dicom_uploader.py index 9c6614e2..65ce5047 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,36 @@ 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_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 "} 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