diff --git a/server.py b/server.py index a7fa6e8..b2605d7 100644 --- a/server.py +++ b/server.py @@ -10,6 +10,7 @@ import requests from mcp.server.fastmcp import FastMCP +from assisted_service_client import models from service_client import InventoryClient from service_client.logger import log @@ -17,6 +18,27 @@ mcp = FastMCP("AssistedService", host="0.0.0.0") +def format_presigned_url(presigned_url: models.PresignedUrl) -> str: + r""" + Format a presigned URL object into a readable string. + + Args: + presigned_url: A PresignedUrl object with url and optional expires_at attributes. + + Returns: + str: A formatted string containing the URL and optional expiration time. + Format: "URL: \nExpires at: " (if expiration exists) + """ + response_parts = [f"URL: {presigned_url.url}"] + # Only include expiration time if it's a meaningful date (not a zero/default value) + if presigned_url.expires_at and not str(presigned_url.expires_at).startswith( + "0001-01-01" + ): + response_parts.append(f"Expires at: {presigned_url.expires_at}") + + return "\n".join(response_parts) + + def get_offline_token() -> str: """ Retrieve the offline token from environment variables or request headers. @@ -228,21 +250,16 @@ async def cluster_iso_download_url(cluster_id: str) -> str: cluster_id, ) - # Extract ISO URLs and expiration dates from each infra env + # Get presigned URLs for each infra env iso_info = [] for infra_env in infra_envs: - iso_url = infra_env.get("download_url") - expires_at = infra_env.get("expires_at") infra_env_id = infra_env.get("id", "unknown") - if iso_url: - # Format the response as a readable string - response_parts = [f"URL: {iso_url}"] - # Only include expiration time if it's a meaningful date (not a zero/default value) - if expires_at and not str(expires_at).startswith("0001-01-01"): - response_parts.append(f"Expires at: {expires_at}") + # Use the new get_infra_env_download_url method + presigned_url = await client.get_infra_env_download_url(infra_env_id) - iso_info.append("\n".join(response_parts)) + if presigned_url.url: + iso_info.append(format_presigned_url(presigned_url)) else: log.warning( "No ISO download URL found for infra env %s", @@ -480,13 +497,7 @@ async def cluster_credentials_download_url(cluster_id: str, file_name: str) -> s result, ) - # Format the response as a readable string - response_parts = [f"URL: {result.url}"] - # Only include expiration time if it's a meaningful date (not a zero/default value) - if result.expires_at and not str(result.expires_at).startswith("0001-01-01"): - response_parts.append(f"Expires at: {result.expires_at}") - - return "\n".join(response_parts) + return format_presigned_url(result) @mcp.tool() diff --git a/service_client/assisted_service_api.py b/service_client/assisted_service_api.py index e6a106e..317fdd6 100644 --- a/service_client/assisted_service_api.py +++ b/service_client/assisted_service_api.py @@ -671,3 +671,46 @@ async def get_presigned_for_cluster_credentials( str(e), ) raise + + async def get_infra_env_download_url( + self, infra_env_id: str + ) -> models.PresignedUrl: + """ + Get presigned download URL for an infrastructure environment. + + Args: + infra_env_id: The unique identifier of the infrastructure environment. + + Returns: + models.PresignedUrl: The presigned URL model containing URL and optional expiration time. + """ + try: + log.info( + "Getting presigned download URL for infrastructure environment %s", + infra_env_id, + ) + result = await asyncio.to_thread( + self._installer_api().get_infra_env_download_url, + infra_env_id=infra_env_id, + ) + log.info( + "Successfully retrieved presigned download URL for infrastructure environment %s", + infra_env_id, + ) + return cast(models.PresignedUrl, result) + except ApiException as e: + log.error( + "API error while getting presigned download URL for infrastructure environment %s: Status: %s, Reason: %s, Body: %s", + infra_env_id, + e.status, + e.reason, + e.body, + ) + raise + except Exception as e: + log.error( + "Unexpected error while getting presigned download URL for infrastructure environment %s: %s", + infra_env_id, + str(e), + ) + raise diff --git a/tests/test_assisted_service_api.py b/tests/test_assisted_service_api.py index 2a63b9f..4fa6fd1 100644 --- a/tests/test_assisted_service_api.py +++ b/tests/test_assisted_service_api.py @@ -685,3 +685,87 @@ async def test_get_presigned_for_cluster_credentials_url_only( mock_api.v2_get_presigned_for_cluster_credentials.assert_called_once_with( cluster_id=cluster_id, file_name=file_name ) + + @pytest.mark.asyncio + async def test_get_infra_env_download_url_success( + self, client: InventoryClient + ) -> None: + """Test successful presigned URL retrieval for infra env download.""" + infra_env_id = "test-infraenv-id" + presigned_url = create_test_presigned_url( + url="https://example.com/infra-env-download", + expires_at="2023-12-31T23:59:59Z", + ) + + with patch.object(client, "_installer_api") as mock_installer_api: + mock_api = Mock() + mock_api.get_infra_env_download_url.return_value = presigned_url + mock_installer_api.return_value = mock_api + + result = await client.get_infra_env_download_url(infra_env_id) + + assert result == presigned_url + mock_api.get_infra_env_download_url.assert_called_once_with( + infra_env_id=infra_env_id + ) + + @pytest.mark.asyncio + async def test_get_infra_env_download_url_api_exception( + self, client: InventoryClient + ) -> None: + """Test infra env download URL retrieval API exception handling.""" + infra_env_id = "test-infraenv-id" + + with patch.object(client, "_installer_api") as mock_installer_api: + mock_api = Mock() + mock_api.get_infra_env_download_url.side_effect = ApiException( + status=404, reason="Not Found" + ) + mock_installer_api.return_value = mock_api + + with pytest.raises(ApiException) as exc_info: + await client.get_infra_env_download_url(infra_env_id) + + assert exc_info.value.status == 404 + assert exc_info.value.reason == "Not Found" + + @pytest.mark.asyncio + async def test_get_infra_env_download_url_unexpected_exception( + self, client: InventoryClient + ) -> None: + """Test infra env download URL retrieval unexpected exception handling.""" + infra_env_id = "test-infraenv-id" + + with patch.object(client, "_installer_api") as mock_installer_api: + mock_api = Mock() + mock_api.get_infra_env_download_url.side_effect = ValueError( + "Unexpected error" + ) + mock_installer_api.return_value = mock_api + + with pytest.raises(ValueError) as exc_info: + await client.get_infra_env_download_url(infra_env_id) + + assert str(exc_info.value) == "Unexpected error" + + @pytest.mark.asyncio + async def test_get_infra_env_download_url_no_expiration( + self, client: InventoryClient + ) -> None: + """Test infra env download URL retrieval when no expiration is returned.""" + infra_env_id = "test-infraenv-id" + presigned_url = create_test_presigned_url( + url="https://example.com/infra-env-download", expires_at=None + ) + + with patch.object(client, "_installer_api") as mock_installer_api: + mock_api = Mock() + mock_api.get_infra_env_download_url.return_value = presigned_url + mock_installer_api.return_value = mock_api + + result = await client.get_infra_env_download_url(infra_env_id) + + assert result == presigned_url + mock_api.get_infra_env_download_url.assert_called_once_with( + infra_env_id=infra_env_id + ) diff --git a/tests/test_server.py b/tests/test_server.py index 3bb9b03..3a27b1c 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -5,7 +5,7 @@ import json import os from typing import Generator, Tuple -from unittest.mock import Mock, patch +from unittest.mock import Mock, patch, call import pytest from requests.exceptions import RequestException @@ -357,10 +357,12 @@ async def test_cluster_iso_download_url_success( "id": "test-infraenv-id", "cluster_id": cluster_id, "openshift_version": "4.18.2", - "download_url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image", - "expires_at": "2023-12-31T23:59:59Z", } mock_inventory_client.list_infra_envs.return_value = [mock_infraenv] + mock_inventory_client.get_infra_env_download_url.return_value = create_test_presigned_url( + url="https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image", + expires_at="2023-12-31T23:59:59Z", + ) with patch.object( server, "InventoryClient", return_value=mock_inventory_client @@ -370,6 +372,9 @@ async def test_cluster_iso_download_url_success( expected_result = "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image\nExpires at: 2023-12-31T23:59:59Z" assert result == expected_result mock_inventory_client.list_infra_envs.assert_called_once_with(cluster_id) + mock_inventory_client.get_infra_env_download_url.assert_called_once_with( + "test-infraenv-id" + ) @pytest.mark.asyncio async def test_cluster_iso_download_url_multiple_infraenvs( @@ -386,8 +391,6 @@ async def test_cluster_iso_download_url_multiple_infraenvs( "id": "test-infraenv-id-1", "cluster_id": cluster_id, "openshift_version": "4.18.2", - "download_url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-1/downloads/image", - "expires_at": "2023-12-31T23:59:59Z", } # Second infraenv with different characteristics @@ -396,8 +399,6 @@ async def test_cluster_iso_download_url_multiple_infraenvs( "id": "test-infraenv-id-2", "cluster_id": cluster_id, "openshift_version": "4.18.2", - "download_url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-2/downloads/image", - "expires_at": "2024-01-15T12:00:00Z", } mock_inventory_client.list_infra_envs.return_value = [ @@ -405,6 +406,18 @@ async def test_cluster_iso_download_url_multiple_infraenvs( mock_infraenv2, ] + # Mock return values for each infra env + mock_inventory_client.get_infra_env_download_url.side_effect = [ + create_test_presigned_url( + url="https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-1/downloads/image", + expires_at="2023-12-31T23:59:59Z", + ), + create_test_presigned_url( + url="https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-2/downloads/image", + expires_at="2024-01-15T12:00:00Z", + ), + ] + with patch.object( server, "InventoryClient", return_value=mock_inventory_client ): @@ -418,6 +431,12 @@ async def test_cluster_iso_download_url_multiple_infraenvs( ) assert result == expected_result mock_inventory_client.list_infra_envs.assert_called_once_with(cluster_id) + mock_inventory_client.get_infra_env_download_url.assert_has_calls( + [ + call("test-infraenv-id-1"), + call("test-infraenv-id-2"), + ] + ) @pytest.mark.asyncio async def test_cluster_iso_download_url_no_expiration( @@ -432,9 +451,12 @@ async def test_cluster_iso_download_url_no_expiration( "id": "test-infraenv-id", "cluster_id": cluster_id, "openshift_version": "4.18.2", - "download_url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image", } mock_inventory_client.list_infra_envs.return_value = [mock_infraenv] + mock_inventory_client.get_infra_env_download_url.return_value = create_test_presigned_url( + url="https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image", + expires_at=None, + ) with patch.object( server, "InventoryClient", return_value=mock_inventory_client @@ -444,6 +466,9 @@ async def test_cluster_iso_download_url_no_expiration( expected_result = "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image" assert result == expected_result mock_inventory_client.list_infra_envs.assert_called_once_with(cluster_id) + mock_inventory_client.get_infra_env_download_url.assert_called_once_with( + "test-infraenv-id" + ) @pytest.mark.asyncio async def test_cluster_iso_download_url_zero_expiration( @@ -458,10 +483,12 @@ async def test_cluster_iso_download_url_zero_expiration( "id": "test-infraenv-id", "cluster_id": cluster_id, "openshift_version": "4.18.2", - "download_url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image", - "expires_at": "0001-01-01 00:00:00+00:00", } mock_inventory_client.list_infra_envs.return_value = [mock_infraenv] + mock_inventory_client.get_infra_env_download_url.return_value = create_test_presigned_url( + url="https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image", + expires_at="0001-01-01 00:00:00+00:00", + ) with patch.object( server, "InventoryClient", return_value=mock_inventory_client @@ -472,6 +499,9 @@ async def test_cluster_iso_download_url_zero_expiration( expected_result = "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image" assert result == expected_result mock_inventory_client.list_infra_envs.assert_called_once_with(cluster_id) + mock_inventory_client.get_infra_env_download_url.assert_called_once_with( + "test-infraenv-id" + ) @pytest.mark.asyncio async def test_cluster_iso_download_url_no_infraenvs(