diff --git a/server.py b/server.py index c8ca85c..009fce0 100644 --- a/server.py +++ b/server.py @@ -8,7 +8,7 @@ import json import os import asyncio -from typing import Optional +from typing import Optional, Any import requests import uvicorn @@ -26,7 +26,7 @@ mcp = FastMCP("AssistedService", host="0.0.0.0", stateless_http=use_stateless_http) -def format_presigned_url(presigned_url: models.PresignedUrl) -> str: +def format_presigned_url(presigned_url: models.PresignedUrl) -> dict[str, Any]: r""" Format a presigned URL object into a readable string. @@ -34,17 +34,26 @@ def format_presigned_url(presigned_url: models.PresignedUrl) -> str: 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}"] + dict: A dict containing URL and optional expiration time. + Format: + { + url: + expires_at: (if expiration exists) + } + """ + presigned_url_dict = { + "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}") + presigned_url_dict["expires_at"] = presigned_url.expires_at.isoformat().replace( + "+00:00", "Z" + ) - return "\n".join(response_parts) + return presigned_url_dict def get_offline_token() -> str: @@ -243,11 +252,12 @@ async def cluster_iso_download_url(cluster_id: str) -> str: cluster_id (str): The unique identifier of the cluster. Returns: - str: A formatted string containing ISO download URLs and optional + dict: A JSON containing ISO download URLs and optional expiration times. Each ISO's information is formatted as: - - URL: - - Expires at: (if available) - Multiple ISOs are separated by blank lines. + [{ + url: + expires_at: (if available) + }] """ log.info("Retrieving InfraEnv ISO URLs for cluster_id: %s", cluster_id) client = InventoryClient(get_access_token()) @@ -287,7 +297,7 @@ async def cluster_iso_download_url(cluster_id: str) -> str: return "No ISO download URLs found for this cluster." log.info("Returning %d ISO URLs for cluster %s", len(iso_info), cluster_id) - return "\n\n".join(iso_info) + return json.dumps(iso_info) @mcp.tool() @@ -537,10 +547,12 @@ async def cluster_credentials_download_url(cluster_id: str, file_name: str) -> s - "kubeadmin-password": The kubeadmin user password file Returns: - str: A formatted string containing the presigned URL and optional + str: A JSON containing the presigned URL and optional expiration time. The response format is: - - URL: - - Expires at: (if available) + { + url: + expires_at: (if available) + } """ log.info( "Getting presigned URL for cluster %s credentials file %s", @@ -556,7 +568,7 @@ async def cluster_credentials_download_url(cluster_id: str, file_name: str) -> s result, ) - return format_presigned_url(result) + return json.dumps(format_presigned_url(result)) @mcp.tool() diff --git a/tests/test_server.py b/tests/test_server.py index 03a2820..5056631 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -371,7 +371,14 @@ async def test_cluster_iso_download_url_success( ): result = await server.cluster_iso_download_url(cluster_id) - expected_result = "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image\nExpires at: 2023-12-31T23:59:59Z" + expected_result = json.dumps( + [ + { + "url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image", + "expires_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( @@ -425,11 +432,17 @@ async def test_cluster_iso_download_url_multiple_infraenvs( ): result = await server.cluster_iso_download_url(cluster_id) - expected_result = ( - "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-1/downloads/image\n" - "Expires at: 2023-12-31T23:59:59Z\n\n" - "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-2/downloads/image\n" - "Expires at: 2024-01-15T12:00:00Z" + expected_result = json.dumps( + [ + { + "url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-1/downloads/image", + "expires_at": "2023-12-31T23:59:59Z", + }, + { + "url": "https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id-2/downloads/image", + "expires_at": "2024-01-15T12:00:00Z", + }, + ] ) assert result == expected_result mock_inventory_client.list_infra_envs.assert_called_once_with(cluster_id) @@ -465,7 +478,13 @@ async def test_cluster_iso_download_url_no_expiration( ): result = await server.cluster_iso_download_url(cluster_id) - expected_result = "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image" + expected_result = json.dumps( + [ + { + "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( @@ -498,7 +517,13 @@ async def test_cluster_iso_download_url_zero_expiration( result = await server.cluster_iso_download_url(cluster_id) # Should not include expiration time since it's a zero/default value - expected_result = "URL: https://api.openshift.com/api/assisted-install/v2/infra-envs/test-id/downloads/image" + expected_result = json.dumps( + [ + { + "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( @@ -826,7 +851,12 @@ async def test_cluster_credentials_download_url_success( cluster_id, file_name ) - expected_result = "URL: https://example.com/presigned-url\nExpires at: 2023-12-31T23:59:59Z" + expected_result = json.dumps( + { + "url": "https://example.com/presigned-url", + "expires_at": "2023-12-31T23:59:59Z", + } + ) assert result == expected_result mock_inventory_client.get_presigned_for_cluster_credentials.assert_called_once_with( cluster_id, file_name @@ -854,7 +884,7 @@ async def test_cluster_credentials_download_url_no_expiration( cluster_id, file_name ) - expected_result = "URL: https://example.com/presigned-url" + expected_result = json.dumps({"url": "https://example.com/presigned-url"}) assert result == expected_result mock_inventory_client.get_presigned_for_cluster_credentials.assert_called_once_with( cluster_id, file_name @@ -885,7 +915,7 @@ async def test_cluster_credentials_download_url_zero_expiration( ) # Should not include expiration time since it's a zero/default value - expected_result = "URL: https://example.com/presigned-url" + expected_result = json.dumps({"url": "https://example.com/presigned-url"}) assert result == expected_result mock_inventory_client.get_presigned_for_cluster_credentials.assert_called_once_with( cluster_id, file_name diff --git a/tests/test_utils.py b/tests/test_utils.py index 07af855..0d73fd1 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -2,6 +2,7 @@ Test utilities for creating test objects. """ +from datetime import datetime from typing import Optional from assisted_service_client import models @@ -84,7 +85,12 @@ def create_test_presigned_url( expires_at: Optional[str] = "2023-12-31T23:59:59Z", ) -> models.PresignedUrl: """Create a test presigned URL object with default values.""" + + date: datetime | None = None + if expires_at is not None: + date = datetime.fromisoformat(expires_at) + return models.PresignedUrl( url=url, - expires_at=expires_at, + expires_at=date, )