diff --git a/src/billing-cost-management-mcp-server/README.md b/src/billing-cost-management-mcp-server/README.md index f69f86238d..e552b9d9ae 100644 --- a/src/billing-cost-management-mcp-server/README.md +++ b/src/billing-cost-management-mcp-server/README.md @@ -39,6 +39,10 @@ MCP server for accessing AWS Billing and Cost Management capabilities. - **Multi-account analysis**: Analyze costs across multiple linked accounts - **Cost driver identification**: Identify key factors driving cost changes +### AWS Billing and Cost Management Pricing Calculator + +- **Workload estimate insights**: Query workload estimates to see what usage you have estimated + ### Specialized Cost Optimization Prompts - **Graviton migration analysis**: Guided analysis to identify EC2 instances suitable for AWS Graviton migration @@ -219,6 +223,12 @@ AWS Pricing: AWS Free Tier: - freetier:GetFreeTierUsage +AWS Billing and Cost Management Pricing Calculator: +- bcm-pricing-calculator:GetPreferences +- bcm-pricing-calculator:GetWorkloadEstimate +- bcm-pricing-calculator:ListWorkloadEstimateUsage +- bcm-pricing-calculator:ListWorkloadEstimates + Storage Lens (Athena and S3): - athena:StartQueryExecution - athena:GetQueryExecution @@ -306,5 +316,11 @@ The server currently supports the following AWS services - get_idle_recommendations - get_enrollment_status -7. **S3 Storage Lens** +7. **Pricing Calculator** + - get-preferences + - get-workload-estimate + - list-workload-estimate-usage + - list-workload-estimates + +8. **S3 Storage Lens** - storage_lens_run_query (custom implementation using Athena) diff --git a/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/server.py b/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/server.py index 77d399e0bb..d11d868380 100755 --- a/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/server.py +++ b/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/server.py @@ -30,6 +30,9 @@ sys.path.insert(0, parent_dir) from awslabs.billing_cost_management_mcp_server.tools.aws_pricing_tools import aws_pricing_server +from awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools import ( + bcm_pricing_calculator_server, +) from awslabs.billing_cost_management_mcp_server.tools.budget_tools import budget_server from awslabs.billing_cost_management_mcp_server.tools.compute_optimizer_tools import ( compute_optimizer_server, @@ -86,6 +89,7 @@ - storage-lens: Query S3 Storage Lens metrics data using Athena SQL - athena-cur: Query Cost and Usage Report data through Athena - pricing: Access AWS service pricing information +- bcm-pricing-calc: Work with workload estimates from AWS Billing and Cost Management Pricing Calculator - budget: Retrieve AWS budget information - cost-anomaly: Identify cost anomalies in AWS accounts - cost-comparison: Compare costs between time periods @@ -137,6 +141,7 @@ async def setup(): await mcp.import_server(cost_optimization_hub_server) await mcp.import_server(storage_lens_server) await mcp.import_server(aws_pricing_server) + await mcp.import_server(bcm_pricing_calculator_server) await mcp.import_server(budget_server) await mcp.import_server(cost_anomaly_server) await mcp.import_server(cost_comparison_server) @@ -157,6 +162,7 @@ async def setup(): 'cost-optimization', 'storage-lens', 'pricing', + 'bcm-pricing-calc', 'budget', 'cost-anomaly', 'cost-comparison', diff --git a/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/tools/bcm_pricing_calculator_tools.py b/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/tools/bcm_pricing_calculator_tools.py new file mode 100644 index 0000000000..946bba2e25 --- /dev/null +++ b/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/tools/bcm_pricing_calculator_tools.py @@ -0,0 +1,899 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""AWS Billing and Cost Management Pricing Calculator tools for the AWS Billing and Cost Management MCP server. + +Updated to use shared utility functions. +""" + +import json +from ..utilities.aws_service_base import ( + create_aws_client, + format_response, + handle_aws_error, + paginate_aws_response, +) +from datetime import datetime +from fastmcp import Context, FastMCP +from typing import Any, Dict, Optional + + +# Constants +DATETIME_FORMAT = '%Y-%m-%d %H:%M:%S UTC' +UTC_TIMEZONE_OFFSET = '+00:00' +BCM_PRICING_CALCULATOR_SERVICE_NAME = 'BCM Pricing Calculator' +PREFERENCES_NOT_CONFIGURED_ERROR = 'BCM Pricing Calculator preferences are not configured. Please configure preferences before using this service.' + +bcm_pricing_calculator_server = FastMCP( + name='bcm-pricing-calc-tools', + instructions=f'{BCM_PRICING_CALCULATOR_SERVICE_NAME} tools for working with AWS Billing and Cost Management Pricing Calculator API', +) + + +async def bcm_pricing_calc_core( + ctx: Context, + operation: str, + identifier: Optional[str] = None, + created_after: Optional[str] = None, + created_before: Optional[str] = None, + expires_after: Optional[str] = None, + expires_before: Optional[str] = None, + status_filter: Optional[str] = None, + name_filter: Optional[str] = None, + name_match_option: str = 'CONTAINS', + usage_account_id_filter: Optional[str] = None, + service_code_filter: Optional[str] = None, + usage_type_filter: Optional[str] = None, + operation_filter: Optional[str] = None, + location_filter: Optional[str] = None, + usage_group_filter: Optional[str] = None, + next_token: Optional[str] = None, + max_results: Optional[int] = None, + max_pages: Optional[int] = None, +) -> Dict[str, Any]: + """Core business logic for BCM Pricing Calculator. + + Args: + ctx: The MCP context object + operation: The operation to perform + identifier: Identifier for specific operations + created_after: Filter estimates created after this timestamp + created_before: Filter estimates created before this timestamp + expires_after: Filter estimates expiring after this timestamp + expires_before: Filter estimates expiring before this timestamp + status_filter: Filter by status + name_filter: Filter by name + name_match_option: Match option for name filter + usage_account_id_filter: Filter by AWS account ID + service_code_filter: Filter by AWS service code + usage_type_filter: Filter by usage type + operation_filter: Filter by operation name + location_filter: Filter by location/region + usage_group_filter: Filter by usage group + next_token: Token for pagination + max_results: Maximum number of results to return + max_pages: Maximum number of API calls to make + + Returns: + Dict containing the response data + """ + try: + # Log the request + await ctx.info(f'Received BCM Pricing Calculator operation: {operation}') + + # Check if the operation is valid + if operation not in [ + 'get_workload_estimate', + 'list_workload_estimates', + 'list_workload_estimate_usage', + 'get_preferences', + ]: + return format_response( + 'error', + {'invalid_parameter': 'operation'}, + f'Invalid operation: {operation}. Valid operations are: get_workload_estimates, get_preferences, describe_workload_estimates', + ) + + # Call the appropriate operation + if operation == 'get_workload_estimate': + return await get_workload_estimate(ctx, identifier) + elif operation == 'list_workload_estimates': + return await list_workload_estimates( + ctx, + created_after, + created_before, + expires_after, + expires_before, + status_filter, + name_filter, + name_match_option, + next_token, + max_results, + max_pages, + ) + elif operation == 'list_workload_estimate_usage': + return await list_workload_estimate_usage( + ctx, + identifier, + usage_account_id_filter, + service_code_filter, + usage_type_filter, + operation_filter, + location_filter, + usage_group_filter, + next_token, + max_results, + max_pages, + ) + elif operation == 'get_preferences': + preferences_result = await get_preferences(ctx) + if 'error' in preferences_result: + return format_response( + 'error', + {'error': preferences_result['error']}, + preferences_result['error'], + ) + else: + return format_response( + 'success', + { + 'message': 'Preferences are properly configured', + 'account_types': preferences_result['account_types'], + }, + ) + else: + return format_response('error', {'message': f'Unknown operation: {operation}'}) + + except Exception as e: + # Use shared error handler for consistent error handling + error_response = await handle_aws_error( + ctx, e, operation, 'AWS Billing and Cost Management Pricing Calculator' + ) + await ctx.error( + f'Failed to process AWS Billing and Cost Management Pricing Calculator request: {error_response.get("data", {}).get("error", str(e))}' + ) + return format_response( + 'error', + {'error': error_response.get('data', {}).get('error', str(e))}, + f'Failed to process AWS Billing and Cost Management Pricing Calculator request: {error_response.get("data", {}).get("error", str(e))}', + ) + + +@bcm_pricing_calculator_server.tool( + name='bcm-pricing-calc', + description="""Allows working with workload estimates using the AWS Billing and Cost Management Pricing Calculator API. + +IMPORTANT USAGE GUIDELINES: +- Always first check the rate preference setting for the authorized principal by calling the get_preferences operation. +- DO NOT state assumptions about Free Tier API + +USE THIS TOOL FOR: +- Listing available **workload estimates** for the logged in account. +- **Filter list of available workload estimates** using name, status, created date, or expiration date. +- Get **details of a workload estimate**. +- Get the list of **services, usage type, operation, and usage amount** modeled within a workload estimate. +- Get **rate preferences** set for Pricing Calculator. These rate preferences denote what rate preferences can be used by each account type in your organization. + +## OPERATIONS + +1) list_workload_estimates - list of available workload estimates + Required: operation="list_workload_estimates" + Optional: created_after, created_before, expires_after, expires_before, status_filter, name_filter, name_match_option, next_token, max_results + Returns: List of all workload estimates for the account. + +2) get_workload_estimate - get details of a workload estimate + Required: operation="get_workload_estimate", identifier + Returns: Details of a specific workload estimate. + +3) list_workload_estimate_usage - list of modeled usage lines within a workload estimate + Required: operation="get_workload_estimate", identifier + Optional: usage_account_id_filter, service_code_filter, usage_type_filter, operation_filter, location_filter, usage_group_filter, next_token, max_results + Returns: List of usage associated with a workload estimate. + +4) get_preferences - get the rate preferences available to an account + Required: operation="get_preferences" + Returns: Retrieves the current preferences for AWS Billing and Cost Management Pricing Calculator. +""", +) +async def bcm_pricing_calc( + ctx: Context, + operation: str, + identifier: Optional[str] = None, + created_after: Optional[str] = None, + created_before: Optional[str] = None, + expires_after: Optional[str] = None, + expires_before: Optional[str] = None, + status_filter: Optional[str] = None, + name_filter: Optional[str] = None, + name_match_option: str = 'CONTAINS', + usage_account_id_filter: Optional[str] = None, + service_code_filter: Optional[str] = None, + usage_type_filter: Optional[str] = None, + operation_filter: Optional[str] = None, + location_filter: Optional[str] = None, + usage_group_filter: Optional[str] = None, + next_token: Optional[str] = None, + max_results: Optional[int] = None, + max_pages: Optional[int] = None, +) -> Dict[str, Any]: + """FastMCP tool wrapper for BCM Pricing Calculator operations.""" + # need this wrapper to improve code coverage as FastMCP decorated methods cannot be tested directly. + return await bcm_pricing_calc_core( + ctx, + operation, + identifier, + created_after, + created_before, + expires_after, + expires_before, + status_filter, + name_filter, + name_match_option, + usage_account_id_filter, + service_code_filter, + usage_type_filter, + operation_filter, + location_filter, + usage_group_filter, + next_token, + max_results, + max_pages, + ) + + +async def get_preferences(ctx: Context) -> dict: + """Check if BCM Pricing Calculator preferences are properly configured. + + Args: + ctx: The MCP context object + + Returns: + dict: Contains either 'account_types' list if preferences are valid, + or 'error' message if not found or error occurred + """ + try: + # Get the BCM Pricing Calculator client + bcm_client = create_aws_client('bcm-pricing-calculator', region_name='us-east-1') + + await ctx.info('Checking BCM Pricing Calculator preferences...') + response = bcm_client.get_preferences() + + # Check if the response contains valid preferences for any account type + if response and ( + 'managementAccountRateTypeSelections' in response + or 'memberAccountRateTypeSelections' in response + or 'standaloneAccountRateTypeSelections' in response + ): + # Log which type of account preferences were found + account_types = [] + if 'managementAccountRateTypeSelections' in response: + account_types.append('management account') + if 'memberAccountRateTypeSelections' in response: + account_types.append('member account') + if 'standaloneAccountRateTypeSelections' in response: + account_types.append('standalone account') + + await ctx.info( + f'BCM Pricing Calculator preferences are properly configured for: {", ".join(account_types)}' + ) + return {'account_types': account_types} + else: + error_msg = 'BCM Pricing Calculator preferences are not configured - no rate type selections found' + await ctx.error(error_msg) + return {'error': error_msg} # the `error` moniker here is used in referenced method. + + except Exception as e: + # Use shared error handler for consistent error handling + error_response = await handle_aws_error( + ctx, e, 'get_preferences', BCM_PRICING_CALCULATOR_SERVICE_NAME + ) + error_msg = f'Failed to check BCM Pricing Calculator preferences: {error_response.get("data", {}).get("error", str(e))}' + await ctx.error(error_msg) + return {'error': error_msg} + + +async def list_workload_estimates( + ctx: Context, + created_after: Optional[str] = None, + created_before: Optional[str] = None, + expires_after: Optional[str] = None, + expires_before: Optional[str] = None, + status_filter: Optional[str] = None, + name_filter: Optional[str] = None, + name_match_option: str = 'CONTAINS', + next_token: Optional[str] = None, + max_results: Optional[int] = None, + max_pages: Optional[int] = None, +) -> Dict[str, Any]: + """Lists all workload estimates for the account. + + Args: + ctx: The MCP context object + created_after: Filter estimates created after this timestamp (ISO format: YYYY-MM-DDTHH:MM:SS) + created_before: Filter estimates created before this timestamp (ISO format: YYYY-MM-DDTHH:MM:SS) + expires_after: Filter estimates expiring after this timestamp (ISO format: YYYY-MM-DDTHH:MM:SS) + expires_before: Filter estimates expiring before this timestamp (ISO format: YYYY-MM-DDTHH:MM:SS) + status_filter: Filter by status (UPDATING, VALID, INVALID, ACTION_NEEDED) + name_filter: Filter by name (supports partial matching) + name_match_option: Match option for name filter (EQUALS, STARTS_WITH, CONTAINS) + next_token: Token for pagination + max_results: Maximum number of results to return + max_pages: Maximum number of API calls to make + + Returns: + Dict containing the workload estimates information. This contains the following information about a workload estimate: + id: The unique identifier of the workload estimate. + name: The name of the workload estimate. + status: The current status of the workload estimate. Possible values are UPDATIN, VALID, INVALID, ACTION_NEEDED + """ + try: + # Log the request + await ctx.info( + f'Listing workload estimates (max_results={max_results}, ' + f'status_filter={status_filter}, name_filter={name_filter})' + ) + + # Create BCM Pricing Calculator client + bcm_client = create_aws_client('bcm-pricing-calculator') + + # Check preferences before proceeding + preferences_result = await get_preferences(ctx) + if 'error' in preferences_result: + return format_response( + 'error', + { + 'error': preferences_result['error'], + 'error_code': 'PREFERENCES_NOT_CONFIGURED', + }, + ) + + request_params: Dict[str, Any] = {} + # Build request parameters + if max_results: + request_params['maxResults'] = max_results + + if next_token: + request_params['nextToken'] = next_token + + # Add created at filter + if created_after or created_before: + created_filter = {} + if created_after: + created_filter['afterTimestamp'] = datetime.fromisoformat( + created_after.replace('Z', UTC_TIMEZONE_OFFSET) + ) + if created_before: + created_filter['beforeTimestamp'] = datetime.fromisoformat( + created_before.replace('Z', UTC_TIMEZONE_OFFSET) + ) + request_params['createdAtFilter'] = created_filter + + # Add expires at filter + if expires_after or expires_before: + expires_filter = {} + if expires_after: + expires_filter['afterTimestamp'] = datetime.fromisoformat( + expires_after.replace('Z', UTC_TIMEZONE_OFFSET) + ) + if expires_before: + expires_filter['beforeTimestamp'] = datetime.fromisoformat( + expires_before.replace('Z', UTC_TIMEZONE_OFFSET) + ) + request_params['expiresAtFilter'] = expires_filter + + # Add additional filters + filters = [] + if status_filter: + filters.append({'name': 'STATUS', 'values': [status_filter], 'matchOption': 'EQUALS'}) + + if name_filter: + filters.append( + {'name': 'NAME', 'values': [name_filter], 'matchOption': name_match_option} + ) + + if filters: + request_params['filters'] = filters + + await ctx.info( + f'Making API call with parameters: {json.dumps(request_params, default=str)}' + ) + + # Handle pagination using shared utility + if max_pages: + # For paginated requests, use the paginate utility + results, pagination_metadata = await paginate_aws_response( + ctx, + 'list_workload_estimates', + lambda **params: bcm_client.list_workload_estimates(**params), + request_params, + 'items', + 'nextToken', + 'nextToken', + max_pages, + ) + + # Format the response + formatted_estimates = [ + format_workload_estimate_response(estimate) for estimate in results + ] + + await ctx.info(f'Retrieved {len(formatted_estimates)} workload estimates') + + # Return success response with pagination metadata + return format_response( + 'success', + { + 'workload_estimates': formatted_estimates, + 'pagination': pagination_metadata, + }, + ) + else: + # For single page, make direct call + response = bcm_client.list_workload_estimates(**request_params) + + # Format the response + formatted_estimates = [ + format_workload_estimate_response(estimate) + for estimate in response.get('items', []) + ] + + await ctx.info(f'Retrieved {len(formatted_estimates)} workload estimates') + + # Return success response using shared format_response utility + return format_response( + 'success', + { + 'workload_estimates': formatted_estimates, + 'total_count': len(formatted_estimates), + 'next_token': response.get('nextToken'), + 'has_more_results': bool(response.get('nextToken')), + }, + ) + + except Exception as e: + # Use shared error handler for all exceptions (ClientError and others) + return await handle_aws_error( + ctx, e, 'list_workload_estimates', BCM_PRICING_CALCULATOR_SERVICE_NAME + ) + + +async def get_workload_estimate( + ctx: Context, + identifier: Optional[str] = None, +) -> Dict[str, Any]: + """Retrieves details of a specific workload estimate using the AWS Billing and Cost Management Pricing Calculator API. + + This tool uses the GetWorkloadEstimate API to retrieve detailed information about a single workload estimate. + + The API returns comprehensive information about: + - Workload estimate ID and name + - Creation and expiration timestamps + - Rate type and timestamp + - Current status of the estimate + - Total estimated cost and currency + - Failure message if applicable + + REQUIRED PARAMETER: + - identifier: The unique identifier of the workload estimate to retrieve + + POSSIBLE STATUSES: + - UPDATING: The estimate is being updated + - VALID: The estimate is valid and up-to-date + - INVALID: The estimate is invalid + - ACTION_NEEDED: User action is required + + The tool provides formatted results with human-readable timestamps and cost information. + """ + try: + # The reason to have the following "unnecessary" check is because how each MCP tool is registered. + # Each MCP tool is registered with a unique name, irrespective of operations it can perform. + # Thereby there is a single entry point that accepts params required across all operations and routes the call flow to an operation. + # So some paramters could required for one operation while not be required for some other operation. + # Thereby all parameters to the entry point are optional, requiring this check. + if identifier is None: + await ctx.error('Identifier is required when calling get_workload_estimate') + return format_response( + 'error', + { + 'error': 'Identifier is required when calling get_workload_estimate', + 'error_code': 'MISSING_PARAMETER', + }, + ) + + # Log the request + await ctx.info(f'Getting workload estimate details for identifier: {identifier}') + + # Create BCM Pricing Calculator client + bcm_client = create_aws_client('bcm-pricing-calculator') + + # Check preferences before proceeding + preferences_result = await get_preferences(ctx) + if 'error' in preferences_result: + return format_response( + 'error', + { + 'error': preferences_result['error'], + 'error_code': 'PREFERENCES_NOT_CONFIGURED', + }, + ) + + # Build request parameters + request_params: Dict[str, Any] = {'identifier': identifier} + + await ctx.info( + f'Making API call with parameters: {json.dumps(request_params, default=str)}' + ) + + # Call the API + response = bcm_client.get_workload_estimate(**request_params) + + # Format the single workload estimate response + formatted_estimate = format_workload_estimate_response(response) + + await ctx.info(f'Retrieved workload estimate: {formatted_estimate.get("name", "Unknown")}') + + # Return success response using shared format_response utility + return format_response( + 'success', + { + 'workload_estimate': formatted_estimate, + 'identifier': identifier, + }, + ) + + except Exception as e: + # Use shared error handler for all exceptions (ClientError and others) + return await handle_aws_error( + ctx, e, 'get_workload_estimate', BCM_PRICING_CALCULATOR_SERVICE_NAME + ) + + +async def list_workload_estimate_usage( + ctx: Context, + workload_estimate_id: Optional[str] = None, + usage_account_id_filter: Optional[str] = None, + service_code_filter: Optional[str] = None, + usage_type_filter: Optional[str] = None, + operation_filter: Optional[str] = None, + location_filter: Optional[str] = None, + usage_group_filter: Optional[str] = None, + next_token: Optional[str] = None, + max_results: Optional[int] = None, + max_pages: Optional[int] = None, +) -> Dict[str, Any]: + """Core business logic for listing usage entries for a specific workload estimate. + + Args: + ctx: The MCP context object + workload_estimate_id: The unique identifier of the workload estimate + usage_account_id_filter: Filter by AWS account ID + service_code_filter: Filter by AWS service code (e.g., AmazonEC2, AmazonS3) + usage_type_filter: Filter by usage type + operation_filter: Filter by operation name + location_filter: Filter by location/region + usage_group_filter: Filter by usage group + next_token: Token for pagination + max_results: Maximum number of results to return + max_pages: Maximum number of API calls to make + + Returns: + Dict containing the workload estimate usage information + """ + try: + # The reason to have the following "unnecessary" check is because how each MCP tool is registered. + # Each MCP tool is registered with a unique name, irrespective of operations it can perform. + # Thereby there is a single entry point that accepts params required across all operations and routes the call flow to an operation. + # So some paramters could required for one operation while not be required for some other operation. + # Thereby all parameters to the entry point are optional, requiring this check. + if workload_estimate_id is None: + await ctx.error( + 'workload_estimate_id is required when calling list_workload_estimate_usage' + ) + return format_response( + 'error', + { + 'error': 'workload_estimate_id is required when calling list_workload_estimate_usage', + 'error_code': 'MISSING_PARAMETER', + }, + ) + + # Log the request + await ctx.info( + f'Listing workload estimate usage (workload_estimate_id={workload_estimate_id}, ' + f'max_results={max_results}, service_code_filter={service_code_filter})' + ) + + # Create BCM Pricing Calculator client + bcm_client = create_aws_client('bcm-pricing-calculator') + + # Check preferences before proceeding + preferences_result = await get_preferences(ctx) + if 'error' in preferences_result: + return format_response( + 'error', + { + 'error': preferences_result['error'], + 'error_code': 'PREFERENCES_NOT_CONFIGURED', + }, + ) + + request_params: Dict[str, Any] = {} + # Build request parameters + request_params['workloadEstimateId'] = workload_estimate_id + + if max_results: + request_params['maxResults'] = max_results + + if next_token: + request_params['nextToken'] = next_token + + # Add filters + filters = [] + if usage_account_id_filter: + filters.append( + { + 'name': 'USAGE_ACCOUNT_ID', + 'values': [usage_account_id_filter], + 'matchOption': 'EQUALS', + } + ) + + if service_code_filter: + filters.append( + {'name': 'SERVICE_CODE', 'values': [service_code_filter], 'matchOption': 'EQUALS'} + ) + + if usage_type_filter: + filters.append( + {'name': 'USAGE_TYPE', 'values': [usage_type_filter], 'matchOption': 'CONTAINS'} + ) + + if operation_filter: + filters.append( + {'name': 'OPERATION', 'values': [operation_filter], 'matchOption': 'CONTAINS'} + ) + + if location_filter: + filters.append( + {'name': 'LOCATION', 'values': [location_filter], 'matchOption': 'EQUALS'} + ) + + if usage_group_filter: + filters.append( + {'name': 'USAGE_GROUP', 'values': [usage_group_filter], 'matchOption': 'EQUALS'} + ) + + if filters: + request_params['filters'] = filters + + await ctx.info( + f'Making API call with parameters: {json.dumps(request_params, default=str)}' + ) + + # Handle pagination using shared utility + if max_pages: + # For paginated requests, use the paginate utility + results, pagination_metadata = await paginate_aws_response( + ctx, + 'list_workload_estimate_usage', + lambda **params: bcm_client.list_workload_estimate_usage(**params), + request_params, + 'items', + 'nextToken', + 'nextToken', + max_pages, + ) + + # Format the response + formatted_usage_items = [format_usage_item_response(item) for item in results] + + await ctx.info(f'Retrieved {len(formatted_usage_items)} usage items') + + # Return success response with pagination metadata + return format_response( + 'success', + { + 'usage_items': formatted_usage_items, + 'pagination': pagination_metadata, + 'workload_estimate_id': workload_estimate_id, + }, + ) + else: + # For single page, make direct call + response = bcm_client.list_workload_estimate_usage(**request_params) + + # Format the response + formatted_usage_items = [ + format_usage_item_response(item) for item in response.get('items', []) + ] + + await ctx.info(f'Retrieved {len(formatted_usage_items)} usage items') + + # Return success response using shared format_response utility + return format_response( + 'success', + { + 'usage_items': formatted_usage_items, + 'total_count': len(formatted_usage_items), + 'next_token': response.get('nextToken'), + 'has_more_results': bool(response.get('nextToken')), + 'workload_estimate_id': workload_estimate_id, + }, + ) + + except Exception as e: + # Use shared error handler for all exceptions (ClientError and others) + return await handle_aws_error( + ctx, e, 'list_workload_estimate_usage', BCM_PRICING_CALCULATOR_SERVICE_NAME + ) + + +def format_usage_item_response(usage_item: Dict[str, Any]) -> Dict[str, Any]: + """Formats a single usage item object from the list_workload_estimate_usage API response. + + Args: + usage_item: Single usage item object from AWS Billing and Cost Management Pricing Calculator. + + Returns: + Formatted usage item object. + """ + formatted_item = { + 'id': usage_item.get('id'), + 'service_code': usage_item.get('serviceCode'), + 'usage_type': usage_item.get('usageType'), + 'operation': usage_item.get('operation'), + 'location': usage_item.get('location'), + 'usage_account_id': usage_item.get('usageAccountId'), + 'group': usage_item.get('group'), + 'status': usage_item.get('status'), + 'currency': usage_item.get('currency', 'USD'), + } + + # Add quantity information + if 'quantity' in usage_item and usage_item['quantity']: + quantity = usage_item['quantity'] + formatted_item['quantity'] = { + 'amount': quantity.get('amount'), + 'unit': quantity.get('unit'), + 'formatted': f'{quantity.get("amount", 0):,.2f} {quantity.get("unit", "")}' + if quantity.get('amount') is not None + else None, + } + + # Add cost information + if 'cost' in usage_item and usage_item['cost'] is not None: + cost = usage_item['cost'] + currency = usage_item.get('currency', 'USD') + formatted_item['cost'] = { + 'amount': cost, + 'currency': currency, + 'formatted': f'{currency} {cost:,.2f}', + } + + # Add historical usage information if present + if 'historicalUsage' in usage_item and usage_item['historicalUsage']: + historical = usage_item['historicalUsage'] + formatted_historical = { + 'service_code': historical.get('serviceCode'), + 'usage_type': historical.get('usageType'), + 'operation': historical.get('operation'), + 'location': historical.get('location'), + 'usage_account_id': historical.get('usageAccountId'), + } + + # Add bill interval if present + if 'billInterval' in historical and historical['billInterval']: + interval = historical['billInterval'] + formatted_historical['bill_interval'] = { + 'start': interval.get('start').isoformat() if interval.get('start') else None, + 'end': interval.get('end').isoformat() if interval.get('end') else None, + } + + formatted_item['historical_usage'] = formatted_historical + + # Add status indicator + status = usage_item.get('status') + if status: + status_indicators = { + 'VALID': 'Valid', + 'INVALID': 'Invalid', + 'STALE': 'Stale', + } + formatted_item['status_indicator'] = status_indicators.get(status, f'❓ {status}') + + return formatted_item + + +def format_workload_estimate_response(estimate: Dict[str, Any]) -> Dict[str, Any]: + """Formats a single workload estimate object from the get_workload_estimate API response. + + Args: + estimate: Single workload estimate object from the AWS API. + + Returns: + Formatted workload estimate object. + """ + formatted_estimate = { + 'id': estimate.get('id'), + 'name': estimate.get('name'), + 'status': estimate.get('status'), + 'rate_type': estimate.get('rateType'), + } + + # Add timestamps with formatting + if 'createdAt' in estimate: + created_at = estimate['createdAt'] + formatted_estimate['created_at'] = { + 'timestamp': created_at.isoformat() + if isinstance(created_at, datetime) + else created_at, + 'formatted': ( + created_at.strftime(DATETIME_FORMAT) + if isinstance(created_at, datetime) + else created_at + ), + } + + if 'expiresAt' in estimate: + expires_at = estimate['expiresAt'] + formatted_estimate['expires_at'] = { + 'timestamp': expires_at.isoformat() + if isinstance(expires_at, datetime) + else expires_at, + 'formatted': ( + expires_at.strftime(DATETIME_FORMAT) + if isinstance(expires_at, datetime) + else expires_at + ), + } + + if 'rateTimestamp' in estimate: + rate_timestamp = estimate['rateTimestamp'] + formatted_estimate['rate_timestamp'] = { + 'timestamp': rate_timestamp.isoformat() + if isinstance(rate_timestamp, datetime) + else rate_timestamp, + 'formatted': ( + rate_timestamp.strftime(DATETIME_FORMAT) + if isinstance(rate_timestamp, datetime) + else rate_timestamp + ), + } + + # Add cost information + if 'totalCost' in estimate: + total_cost = estimate['totalCost'] + cost_currency = estimate.get('costCurrency', 'USD') + formatted_estimate['cost'] = { + 'amount': total_cost, + 'currency': cost_currency, + 'formatted': f'{cost_currency} {total_cost:,.2f}' if total_cost is not None else None, + } + + # Add failure message if present + if 'failureMessage' in estimate and estimate['failureMessage']: + formatted_estimate['failure_message'] = estimate['failureMessage'] + + # Add status indicator + status = estimate.get('status') + if status: + status_indicators = { + 'VALID': 'Valid', + 'UPDATING': 'Updating', + 'INVALID': 'Invalid', + 'ACTION_NEEDED': 'Action Needed', + } + formatted_estimate['status_indicator'] = status_indicators.get(status, f'❓ {status}') + + return formatted_estimate diff --git a/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/utilities/aws_service_base.py b/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/utilities/aws_service_base.py index 6589b3182a..8cf0dbbe06 100644 --- a/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/utilities/aws_service_base.py +++ b/src/billing-cost-management-mcp-server/awslabs/billing_cost_management_mcp_server/utilities/aws_service_base.py @@ -109,6 +109,7 @@ def create_aws_client(service_name: str, region_name: Optional[str] = None) -> A 'sts', # STS (for account validation) 'freetier', # AWS Free Tier Usage 's3', # AWS S3 + 'bcm-pricing-calculator', # BCM Pricing Calculator ] # Validate requested service diff --git a/src/billing-cost-management-mcp-server/tests/tools/test_aws_bcm_pricing_calculator_tools.py b/src/billing-cost-management-mcp-server/tests/tools/test_aws_bcm_pricing_calculator_tools.py new file mode 100644 index 0000000000..574539968e --- /dev/null +++ b/src/billing-cost-management-mcp-server/tests/tools/test_aws_bcm_pricing_calculator_tools.py @@ -0,0 +1,1895 @@ +"""Tests for BCM Pricing Calculator Tools. + +This module contains comprehensive tests for all methods in the BCM Pricing Calculator Tools, +ensuring complete code coverage and proper error handling. +""" + +import pytest +from awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools import ( + PREFERENCES_NOT_CONFIGURED_ERROR, + bcm_pricing_calc_core, + bcm_pricing_calculator_server, + format_usage_item_response, + format_workload_estimate_response, + get_preferences, + get_workload_estimate, + list_workload_estimate_usage, + list_workload_estimates, +) +from botocore.exceptions import BotoCoreError, ClientError +from datetime import datetime +from fastmcp import Context +from unittest.mock import AsyncMock, MagicMock, patch + + +@pytest.fixture +def mock_context(): + """Create a mock MCP context.""" + context = MagicMock(spec=Context) + context.info = AsyncMock() + context.error = AsyncMock() + return context + + +@pytest.fixture +def mock_bcm_pricing_calculator_client(): + """Create a mock BCM Pricing Calculator boto3 client.""" + mock_client = MagicMock() + + # Set up mock responses for different operations + mock_client.get_preferences.return_value = { + 'managementAccountRateTypeSelections': ['BEFORE_DISCOUNTS'], + 'memberAccountRateTypeSelections': ['BEFORE_DISCOUNTS'], + 'standaloneAccountRateTypeSelections': ['BEFORE_DISCOUNTS'], + } + + mock_client.list_workload_estimates.return_value = { + 'items': [ + { + 'id': 'estimate-123', + 'name': 'Test Workload Estimate', + 'status': 'VALID', + 'rateType': 'BEFORE_DISCOUNTS', + 'createdAt': datetime(2023, 1, 1, 12, 0, 0), + 'expiresAt': datetime(2023, 12, 31, 23, 59, 59), + 'rateTimestamp': datetime(2023, 1, 1, 0, 0, 0), + 'totalCost': 1500.50, + 'costCurrency': 'USD', + }, + { + 'id': 'estimate-456', + 'name': 'Another Estimate', + 'status': 'UPDATING', + 'rateType': 'AFTER_DISCOUNTS', + 'createdAt': datetime(2023, 2, 1, 10, 0, 0), + 'expiresAt': datetime(2023, 12, 31, 23, 59, 59), + 'rateTimestamp': datetime(2023, 2, 1, 0, 0, 0), + 'totalCost': 2000.75, + 'costCurrency': 'USD', + }, + ], + 'nextToken': None, + } + + mock_client.get_workload_estimate.return_value = { + 'id': 'estimate-123', + 'name': 'Test Workload Estimate', + 'status': 'VALID', + 'rateType': 'BEFORE_DISCOUNTS', + 'createdAt': datetime(2023, 1, 1, 12, 0, 0), + 'expiresAt': datetime(2023, 12, 31, 23, 59, 59), + 'rateTimestamp': datetime(2023, 1, 1, 0, 0, 0), + 'totalCost': 1500.50, + 'costCurrency': 'USD', + } + + mock_client.list_workload_estimate_usage.return_value = { + 'items': [ + { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'usageType': 'BoxUsage:t3.medium', + 'operation': 'RunInstances', + 'location': 'US East (N. Virginia)', + 'usageAccountId': '123456789012', + 'group': 'EC2-Instance', + 'status': 'VALID', + 'currency': 'USD', + 'quantity': { + 'amount': 744.0, + 'unit': 'Hrs', + }, + 'cost': 50.25, + 'historicalUsage': { + 'serviceCode': 'AmazonEC2', + 'usageType': 'BoxUsage:t3.medium', + 'operation': 'RunInstances', + 'location': 'US East (N. Virginia)', + 'usageAccountId': '123456789012', + 'billInterval': { + 'start': datetime(2023, 1, 1), + 'end': datetime(2023, 1, 31), + }, + }, + }, + ], + 'nextToken': None, + } + + return mock_client + + +@pytest.mark.asyncio +class TestGetPreferences: + """Tests for get_preferences function.""" + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_success_management_account( + self, mock_create_client, mock_context, mock_bcm_pricing_calculator_client + ): + """Test get_preferences returns dict with account_types when management account preferences are configured.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_bcm_pricing_calculator_client.get_preferences.return_value = { + 'managementAccountRateTypeSelections': ['BEFORE_DISCOUNTS'] + } + + # Execute + result = await get_preferences(mock_context) + + # Assert + mock_create_client.assert_called_once_with( + 'bcm-pricing-calculator', region_name='us-east-1' + ) + mock_bcm_pricing_calculator_client.get_preferences.assert_called_once() + assert 'account_types' in result + assert 'management account' in result['account_types'] + mock_context.info.assert_called() + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_success_member_account( + self, mock_create_client, mock_context, mock_bcm_pricing_calculator_client + ): + """Test get_preferences returns dict with account_types when member account preferences are configured.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_bcm_pricing_calculator_client.get_preferences.return_value = { + 'memberAccountRateTypeSelections': ['AFTER_DISCOUNTS'] + } + + # Execute + result = await get_preferences(mock_context) + + # Assert + assert 'account_types' in result + assert 'member account' in result['account_types'] + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_success_standalone_account( + self, mock_create_client, mock_context, mock_bcm_pricing_calculator_client + ): + """Test get_preferences returns dict with account_types when standalone account preferences are configured.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_bcm_pricing_calculator_client.get_preferences.return_value = { + 'standaloneAccountRateTypeSelections': ['AFTER_DISCOUNTS_AND_COMMITMENTS'] + } + + # Execute + result = await get_preferences(mock_context) + + # Assert + assert 'account_types' in result + assert 'standalone account' in result['account_types'] + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_not_configured( + self, mock_create_client, mock_context, mock_bcm_pricing_calculator_client + ): + """Test get_preferences returns dict with error when no preferences are configured.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_bcm_pricing_calculator_client.get_preferences.return_value = {} + + # Execute + result = await get_preferences(mock_context) + + # Assert + assert 'error' in result + assert 'BCM Pricing Calculator preferences are not configured' in result['error'] + mock_context.error.assert_called() + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.handle_aws_error' + ) + async def test_get_preferences_exception( + self, mock_handle_error, mock_create_client, mock_context + ): + """Test get_preferences handles exceptions properly.""" + # Setup + error = ClientError( + {'Error': {'Code': 'AccessDenied', 'Message': 'Access denied'}}, 'GetPreferences' + ) + mock_create_client.side_effect = error + mock_handle_error.return_value = {'data': {'error': 'Access denied'}} + + # Execute + result = await get_preferences(mock_context) + + # Assert + mock_handle_error.assert_called_once_with( + mock_context, error, 'get_preferences', 'BCM Pricing Calculator' + ) + assert 'error' in result + assert ( + 'Failed to check BCM Pricing Calculator preferences: Access denied' in result['error'] + ) + mock_context.error.assert_called() + + +@pytest.mark.asyncio +class TestListWorkloadEstimates: + """Tests for list_workload_estimates function.""" + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_success( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates returns formatted estimates.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': 'management account'} + + # Execute + result = await list_workload_estimates(mock_context, max_results=50) + + # Assert + mock_create_client.assert_called_once_with('bcm-pricing-calculator') + mock_get_preferences.assert_called_once() + mock_bcm_pricing_calculator_client.list_workload_estimates.assert_called_once() + + assert result['status'] == 'success' + assert 'workload_estimates' in result['data'] + assert len(result['data']['workload_estimates']) == 2 + assert result['data']['total_count'] == 2 + assert result['data']['has_more_results'] is False + + # Check first estimate details + first_estimate = result['data']['workload_estimates'][0] + assert first_estimate['id'] == 'estimate-123' + assert first_estimate['name'] == 'Test Workload Estimate' + assert first_estimate['status'] == 'VALID' + assert first_estimate['status_indicator'] == 'Valid' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_with_filters( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates with various filters.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': 'management account'} + + # Execute + result = await list_workload_estimates( + mock_context, + created_after='2023-01-01T00:00:00Z', + created_before='2023-12-31T23:59:59Z', + expires_after='2023-06-01T00:00:00Z', + expires_before='2024-01-01T00:00:00Z', + status_filter='VALID', + name_filter='Test', + name_match_option='CONTAINS', + max_results=25, + ) + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + assert call_kwargs['maxResults'] == 25 + assert 'createdAtFilter' in call_kwargs + assert 'expiresAtFilter' in call_kwargs + assert 'filters' in call_kwargs + assert len(call_kwargs['filters']) == 2 # status and name filters + + assert result['status'] == 'success' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + async def test_list_workload_estimates_preferences_not_configured( + self, mock_get_preferences, mock_context + ): + """Test list_workload_estimates when preferences are not configured.""" + # Setup + mock_get_preferences.return_value = { + 'error': 'BCM Pricing Calculator preferences are not configured' + } + + # Execute + result = await list_workload_estimates(mock_context) + + # Assert + assert result['status'] == 'error' + assert result['data']['error_code'] == 'PREFERENCES_NOT_CONFIGURED' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.handle_aws_error' + ) + async def test_list_workload_estimates_exception( + self, mock_handle_error, mock_create_client, mock_get_preferences, mock_context + ): + """Test list_workload_estimates handles exceptions properly.""" + # Setup + error = ClientError( + {'Error': {'Code': 'ValidationException', 'Message': 'Invalid parameter'}}, + 'ListWorkloadEstimates', + ) + mock_get_preferences.return_value = {'account_types': 'management account'} + mock_create_client.return_value.list_workload_estimates.side_effect = error + mock_handle_error.return_value = {'status': 'error', 'message': 'Invalid parameter'} + + # Execute + result = await list_workload_estimates(mock_context) + + # Assert + mock_handle_error.assert_called_once_with( + mock_context, error, 'list_workload_estimates', 'BCM Pricing Calculator' + ) + assert result['status'] == 'error' + + +@pytest.mark.asyncio +class TestGetWorkloadEstimate: + """Tests for get_workload_estimate function.""" + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_workload_estimate_success( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test get_workload_estimate returns formatted estimate.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': 'management account'} + + # Execute + result = await get_workload_estimate(mock_context, identifier='estimate-123') + + # Assert + mock_create_client.assert_called_once_with('bcm-pricing-calculator') + mock_get_preferences.assert_called_once() + mock_bcm_pricing_calculator_client.get_workload_estimate.assert_called_once_with( + identifier='estimate-123' + ) + + assert result['status'] == 'success' + assert 'workload_estimate' in result['data'] + assert result['data']['identifier'] == 'estimate-123' + + estimate = result['data']['workload_estimate'] + assert estimate['id'] == 'estimate-123' + assert estimate['name'] == 'Test Workload Estimate' + assert estimate['status'] == 'VALID' + + async def test_get_workload_estimate_missing_identifier(self, mock_context): + """Test get_workload_estimate returns error when identifier is missing.""" + # Execute + result = await get_workload_estimate(mock_context, identifier=None) + + # Assert + assert result['status'] == 'error' + assert 'Identifier is required' in result['data']['error'] + assert result['data']['error_code'] == 'MISSING_PARAMETER' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + async def test_get_workload_estimate_preferences_not_configured( + self, mock_get_preferences, mock_context + ): + """Test get_workload_estimate when preferences are not configured.""" + # Setup + mock_get_preferences.return_value = { + 'error': 'BCM Pricing Calculator preferences are not configured' + } + + # Execute + result = await get_workload_estimate(mock_context, identifier='estimate-123') + + # Assert + assert result['status'] == 'error' + assert result['data']['error_code'] == 'PREFERENCES_NOT_CONFIGURED' + + +@pytest.mark.asyncio +class TestListWorkloadEstimateUsage: + """Tests for list_workload_estimate_usage function.""" + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimate_usage_success( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimate_usage returns formatted usage items.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute + result = await list_workload_estimate_usage( + mock_context, + workload_estimate_id='estimate-123', + service_code_filter='AmazonEC2', + max_results=50, + ) + + # Assert + mock_create_client.assert_called_once_with('bcm-pricing-calculator') + mock_get_preferences.assert_called_once() + mock_bcm_pricing_calculator_client.list_workload_estimate_usage.assert_called_once() + + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimate_usage.call_args[1] + assert call_kwargs['workloadEstimateId'] == 'estimate-123' + assert call_kwargs['maxResults'] == 50 + assert 'filters' in call_kwargs + assert len(call_kwargs['filters']) == 1 # service_code_filter + + assert result['status'] == 'success' + assert 'usage_items' in result['data'] + assert len(result['data']['usage_items']) == 1 + assert result['data']['workload_estimate_id'] == 'estimate-123' + + # Check usage item details + usage_item = result['data']['usage_items'][0] + assert usage_item['id'] == 'usage-123' + assert usage_item['service_code'] == 'AmazonEC2' + assert usage_item['usage_type'] == 'BoxUsage:t3.medium' + + async def test_list_workload_estimate_usage_missing_id(self, mock_context): + """Test list_workload_estimate_usage returns error when workload_estimate_id is missing.""" + # Execute + result = await list_workload_estimate_usage(mock_context, workload_estimate_id=None) + + # Assert + assert result['status'] == 'error' + assert 'workload_estimate_id is required' in result['data']['error'] + assert result['data']['error_code'] == 'MISSING_PARAMETER' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimate_usage_with_all_filters( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimate_usage with all possible filters.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute + result = await list_workload_estimate_usage( + mock_context, + workload_estimate_id='estimate-123', + usage_account_id_filter='123456789012', + service_code_filter='AmazonEC2', + usage_type_filter='BoxUsage', + operation_filter='RunInstances', + location_filter='US East (N. Virginia)', + usage_group_filter='EC2-Instance', + max_results=100, + ) + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimate_usage.call_args[1] + assert len(call_kwargs['filters']) == 6 # All filters applied + assert result['status'] == 'success' + + +class TestFormatWorkloadEstimateResponse: + """Tests for format_workload_estimate_response function.""" + + def test_format_workload_estimate_response_basic(self): + """Test format_workload_estimate_response with basic fields.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'VALID', + 'rateType': 'BEFORE_DISCOUNTS', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert result['id'] == 'estimate-123' + assert result['name'] == 'Test Estimate' + assert result['status'] == 'VALID' + assert result['rate_type'] == 'BEFORE_DISCOUNTS' + assert result['status_indicator'] == 'Valid' + + def test_format_workload_estimate_response_with_timestamps(self): + """Test format_workload_estimate_response with timestamp fields.""" + # Setup + created_at = datetime(2023, 1, 1, 12, 0, 0) + expires_at = datetime(2023, 12, 31, 23, 59, 59) + rate_timestamp = datetime(2023, 1, 1, 0, 0, 0) + + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'UPDATING', + 'createdAt': created_at, + 'expiresAt': expires_at, + 'rateTimestamp': rate_timestamp, + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert 'created_at' in result + assert result['created_at']['timestamp'] == created_at.isoformat() + assert result['created_at']['formatted'] == '2023-01-01 12:00:00 UTC' + + assert 'expires_at' in result + assert result['expires_at']['timestamp'] == expires_at.isoformat() + + assert 'rate_timestamp' in result + assert result['rate_timestamp']['timestamp'] == rate_timestamp.isoformat() + + assert result['status_indicator'] == 'Updating' + + def test_format_workload_estimate_response_with_cost(self): + """Test format_workload_estimate_response with cost information.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'VALID', + 'totalCost': 1500.50, + 'costCurrency': 'USD', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert 'cost' in result + assert result['cost']['amount'] == 1500.50 + assert result['cost']['currency'] == 'USD' + assert result['cost']['formatted'] == 'USD 1,500.50' + + def test_format_workload_estimate_response_with_failure_message(self): + """Test format_workload_estimate_response with failure message.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'INVALID', + 'failureMessage': 'Invalid configuration detected', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert result['failure_message'] == 'Invalid configuration detected' + assert result['status_indicator'] == 'Invalid' + + def test_format_workload_estimate_response_action_needed_status(self): + """Test format_workload_estimate_response with ACTION_NEEDED status.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'ACTION_NEEDED', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert result['status_indicator'] == 'Action Needed' + + def test_format_workload_estimate_response_unknown_status(self): + """Test format_workload_estimate_response with unknown status.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'UNKNOWN_STATUS', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert result['status_indicator'] == '❓ UNKNOWN_STATUS' + + +class TestFormatUsageItemResponse: + """Tests for format_usage_item_response function.""" + + def test_format_usage_item_response_basic(self): + """Test format_usage_item_response with basic fields.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'usageType': 'BoxUsage:t3.medium', + 'operation': 'RunInstances', + 'location': 'US East (N. Virginia)', + 'usageAccountId': '123456789012', + 'group': 'EC2-Instance', + 'status': 'VALID', + 'currency': 'USD', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert result['id'] == 'usage-123' + assert result['service_code'] == 'AmazonEC2' + assert result['usage_type'] == 'BoxUsage:t3.medium' + assert result['operation'] == 'RunInstances' + assert result['location'] == 'US East (N. Virginia)' + assert result['usage_account_id'] == '123456789012' + assert result['group'] == 'EC2-Instance' + assert result['status'] == 'VALID' + assert result['currency'] == 'USD' + assert result['status_indicator'] == 'Valid' + + def test_format_usage_item_response_with_quantity(self): + """Test format_usage_item_response with quantity information.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'quantity': { + 'amount': 744.0, + 'unit': 'Hrs', + }, + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'quantity' in result + assert result['quantity']['amount'] == 744.0 + assert result['quantity']['unit'] == 'Hrs' + assert result['quantity']['formatted'] == '744.00 Hrs' + + def test_format_usage_item_response_with_cost(self): + """Test format_usage_item_response with cost information.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'cost': 50.25, + 'currency': 'USD', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'cost' in result + assert result['cost']['amount'] == 50.25 + assert result['cost']['currency'] == 'USD' + assert result['cost']['formatted'] == 'USD 50.25' + + def test_format_usage_item_response_with_historical_usage(self): + """Test format_usage_item_response with historical usage information.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'historicalUsage': { + 'serviceCode': 'AmazonEC2', + 'usageType': 'BoxUsage:t3.medium', + 'operation': 'RunInstances', + 'location': 'US East (N. Virginia)', + 'usageAccountId': '123456789012', + 'billInterval': { + 'start': datetime(2023, 1, 1), + 'end': datetime(2023, 1, 31), + }, + }, + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'historical_usage' in result + historical = result['historical_usage'] + assert historical['service_code'] == 'AmazonEC2' + assert historical['usage_type'] == 'BoxUsage:t3.medium' + assert historical['operation'] == 'RunInstances' + assert historical['location'] == 'US East (N. Virginia)' + assert historical['usage_account_id'] == '123456789012' + assert 'bill_interval' in historical + assert historical['bill_interval']['start'] == '2023-01-01T00:00:00' + assert historical['bill_interval']['end'] == '2023-01-31T00:00:00' + + def test_format_usage_item_response_quantity_none_amount(self): + """Test format_usage_item_response with None quantity amount.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'quantity': { + 'amount': None, + 'unit': 'Hrs', + }, + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'quantity' in result + assert result['quantity']['amount'] is None + assert result['quantity']['unit'] == 'Hrs' + assert result['quantity']['formatted'] is None + + def test_format_usage_item_response_stale_status(self): + """Test format_usage_item_response with STALE status.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'status': 'STALE', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert result['status_indicator'] == 'Stale' + + def test_format_usage_item_response_invalid_status(self): + """Test format_usage_item_response with INVALID status.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'status': 'INVALID', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert result['status_indicator'] == 'Invalid' + + def test_format_usage_item_response_unknown_status(self): + """Test format_usage_item_response with unknown status.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'status': 'UNKNOWN_STATUS', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert result['status_indicator'] == '❓ UNKNOWN_STATUS' + + +def test_bcm_pricing_calculator_server_initialization(): + """Test that the bcm_pricing_calculator_server is properly initialized.""" + # Verify the server name + assert bcm_pricing_calculator_server.name == 'bcm-pricing-calc-tools' + + # Verify the server instructions + instructions = bcm_pricing_calculator_server.instructions + assert instructions is not None + assert 'BCM Pricing Calculator tools' in instructions + + +class TestAdditionalFormattingCases: + """Tests for additional formatting edge cases to achieve complete coverage.""" + + def test_format_usage_item_response_no_quantity(self): + """Test format_usage_item_response with no quantity field.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'status': 'VALID', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'quantity' not in result + assert result['id'] == 'usage-123' + assert result['service_code'] == 'AmazonEC2' + assert result['status'] == 'VALID' + + def test_format_usage_item_response_no_cost(self): + """Test format_usage_item_response with no cost field.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'status': 'VALID', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'cost' not in result + assert result['id'] == 'usage-123' + assert result['service_code'] == 'AmazonEC2' + + def test_format_usage_item_response_no_status(self): + """Test format_usage_item_response with no status field.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'status_indicator' not in result + assert result['id'] == 'usage-123' + assert result['service_code'] == 'AmazonEC2' + + def test_format_usage_item_response_default_currency(self): + """Test format_usage_item_response uses default USD currency.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + # No currency field provided + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert result['currency'] == 'USD' + + def test_format_usage_item_response_historical_usage_no_bill_interval_start_end(self): + """Test format_usage_item_response with historical usage but no start/end in bill interval.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'historicalUsage': { + 'serviceCode': 'AmazonEC2', + 'usageType': 'BoxUsage:t3.medium', + 'operation': 'RunInstances', + 'location': 'US East (N. Virginia)', + 'usageAccountId': '123456789012', + 'billInterval': { + 'start': None, + 'end': None, + }, + }, + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'historical_usage' in result + historical = result['historical_usage'] + assert 'bill_interval' in historical + assert historical['bill_interval']['start'] is None + assert historical['bill_interval']['end'] is None + + def test_format_workload_estimate_response_no_status(self): + """Test format_workload_estimate_response with no status field.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'rateType': 'BEFORE_DISCOUNTS', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert 'status_indicator' not in result + assert result['id'] == 'estimate-123' + assert result['name'] == 'Test Estimate' + + def test_format_workload_estimate_response_no_failure_message(self): + """Test format_workload_estimate_response with no failure message.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'VALID', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert 'failure_message' not in result + assert result['id'] == 'estimate-123' + assert result['status'] == 'VALID' + + +@pytest.mark.asyncio +class TestMissingCoverageBranches: + """Tests specifically targeting lines 100-148 that are missing coverage.""" + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_management_account_only(self, mock_create_client, mock_context): + """Test get_preferences with only management account preferences (covers lines ~100-110).""" + # Setup - only management account preferences + mock_client = MagicMock() + mock_client.get_preferences.return_value = { + 'managementAccountRateTypeSelections': [ + 'BEFORE_DISCOUNTS', + 'AFTER_DISCOUNTS_AND_COMMITMENTS', + ] + # No member or standalone account selections + } + mock_create_client.return_value = mock_client + + # Execute + result = await get_preferences(mock_context) + + # Assert + assert 'account_types' in result + assert 'management account' in result['account_types'] + mock_context.info.assert_called() + # Verify the specific log message for management account + info_calls = [call.args[0] for call in mock_context.info.call_args_list] + assert any('management account' in call for call in info_calls) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_member_account_only(self, mock_create_client, mock_context): + """Test get_preferences with only member account preferences (covers lines ~100-110).""" + # Setup - only member account preferences + mock_client = MagicMock() + mock_client.get_preferences.return_value = { + 'memberAccountRateTypeSelections': ['BEFORE_DISCOUNTS'] + # No management or standalone account selections + } + mock_create_client.return_value = mock_client + + # Execute + result = await get_preferences(mock_context) + + # Assert + assert 'account_types' in result + assert 'member account' in result['account_types'] + mock_context.info.assert_called() + # Verify the specific log message for member account + info_calls = [call.args[0] for call in mock_context.info.call_args_list] + assert any('member account' in call for call in info_calls) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_standalone_account_only(self, mock_create_client, mock_context): + """Test get_preferences with only standalone account preferences (covers lines ~100-110).""" + # Setup - only standalone account preferences + mock_client = MagicMock() + mock_client.get_preferences.return_value = { + 'standaloneAccountRateTypeSelections': ['AFTER_DISCOUNTS_AND_COMMITMENTS'] + # No management or member account selections + } + mock_create_client.return_value = mock_client + + # Execute + result = await get_preferences(mock_context) + + # Assert + assert 'account_types' in result + assert 'standalone account' in result['account_types'] + mock_context.info.assert_called() + # Verify the specific log message for standalone account + info_calls = [call.args[0] for call in mock_context.info.call_args_list] + assert any('standalone account' in call for call in info_calls) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_get_preferences_multiple_account_types(self, mock_create_client, mock_context): + """Test get_preferences with multiple account type preferences (covers lines ~100-110).""" + # Setup - multiple account type preferences + mock_client = MagicMock() + mock_client.get_preferences.return_value = { + 'managementAccountRateTypeSelections': ['BEFORE_DISCOUNTS'], + 'memberAccountRateTypeSelections': ['AFTER_DISCOUNTS'], + 'standaloneAccountRateTypeSelections': ['AFTER_DISCOUNTS_AND_COMMITMENTS'], + } + mock_create_client.return_value = mock_client + + # Execute + result = await get_preferences(mock_context) + + # Assert + assert 'account_types' in result + assert 'management account' in result['account_types'] + assert 'member account' in result['account_types'] + assert 'standalone account' in result['account_types'] + mock_context.info.assert_called() + # Verify the log message contains all account types + info_calls = [call.args[0] for call in mock_context.info.call_args_list] + combined_message = ' '.join(info_calls) + assert 'management account' in combined_message + assert 'member account' in combined_message + assert 'standalone account' in combined_message + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_datetime_parsing_created_after_only( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates datetime parsing for created_after only (covers lines ~120-130).""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute with only created_after + result = await list_workload_estimates(mock_context, created_after='2023-01-01T00:00:00Z') + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + assert 'createdAtFilter' in call_kwargs + created_filter = call_kwargs['createdAtFilter'] + assert 'afterTimestamp' in created_filter + assert 'beforeTimestamp' not in created_filter + assert result['status'] == 'success' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_datetime_parsing_created_before_only( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates datetime parsing for created_before only (covers lines ~120-130).""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute with only created_before + result = await list_workload_estimates(mock_context, created_before='2023-12-31T23:59:59Z') + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + assert 'createdAtFilter' in call_kwargs + created_filter = call_kwargs['createdAtFilter'] + assert 'beforeTimestamp' in created_filter + assert 'afterTimestamp' not in created_filter + assert result['status'] == 'success' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_datetime_parsing_expires_after_only( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates datetime parsing for expires_after only (covers lines ~130-140).""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute with only expires_after + result = await list_workload_estimates(mock_context, expires_after='2023-06-01T00:00:00Z') + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + assert 'expiresAtFilter' in call_kwargs + expires_filter = call_kwargs['expiresAtFilter'] + assert 'afterTimestamp' in expires_filter + assert 'beforeTimestamp' not in expires_filter + assert result['status'] == 'success' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_datetime_parsing_expires_before_only( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates datetime parsing for expires_before only (covers lines ~130-140).""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute with only expires_before + result = await list_workload_estimates(mock_context, expires_before='2024-01-01T00:00:00Z') + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + assert 'expiresAtFilter' in call_kwargs + expires_filter = call_kwargs['expiresAtFilter'] + assert 'beforeTimestamp' in expires_filter + assert 'afterTimestamp' not in expires_filter + assert result['status'] == 'success' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_filter_building_status_only( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates filter building for status only (covers lines ~140-148).""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute with only status filter + result = await list_workload_estimates(mock_context, status_filter='VALID') + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + assert 'filters' in call_kwargs + filters = call_kwargs['filters'] + assert len(filters) == 1 + assert filters[0]['name'] == 'STATUS' + assert filters[0]['values'] == ['VALID'] + assert filters[0]['matchOption'] == 'EQUALS' + assert result['status'] == 'success' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_filter_building_name_only( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates filter building for name only (covers lines ~140-148).""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute with only name filter + result = await list_workload_estimates( + mock_context, name_filter='Test', name_match_option='STARTS_WITH' + ) + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + assert 'filters' in call_kwargs + filters = call_kwargs['filters'] + assert len(filters) == 1 + assert filters[0]['name'] == 'NAME' + assert filters[0]['values'] == ['Test'] + assert filters[0]['matchOption'] == 'STARTS_WITH' + assert result['status'] == 'success' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + async def test_list_workload_estimates_no_filters_applied( + self, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates with no filters (covers the filters conditional logic).""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute with no filters + result = await list_workload_estimates(mock_context) + + # Assert + call_kwargs = mock_bcm_pricing_calculator_client.list_workload_estimates.call_args[1] + # Should not have filters key when no filters are applied + assert 'filters' not in call_kwargs or not call_kwargs.get('filters') + assert result['status'] == 'success' + + +@pytest.mark.asyncio +class TestAdditionalConditionalBranches: + """Tests for additional conditional branches to achieve complete coverage.""" + + +@pytest.mark.asyncio +class TestListWorkloadEstimateUsagePreferencesNotConfigured: + """Test for line 476 - preferences not configured in list_workload_estimate_usage.""" + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + async def test_list_workload_estimate_usage_preferences_not_configured( + self, mock_get_preferences, mock_context + ): + """Test list_workload_estimate_usage when preferences are not configured (line 476).""" + # Setup + mock_get_preferences.return_value = { + 'error': 'BCM Pricing Calculator preferences are not configured - no rate type selections found' + } + + # Execute + result = await list_workload_estimate_usage( + mock_context, workload_estimate_id='estimate-123' + ) + + # Assert + assert result['status'] == 'error' + assert ( + result['data']['error'] + == 'BCM Pricing Calculator preferences are not configured - no rate type selections found' + ) + assert result['data']['error_code'] == 'PREFERENCES_NOT_CONFIGURED' + + +class TestEdgeCases: + """Tests for edge cases and error conditions.""" + + @pytest.mark.asyncio + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.paginate_aws_response' + ) + async def test_list_workload_estimates_with_pagination( + self, + mock_paginate_aws_response, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimates with pagination token.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Mock the paginate_aws_response function to prevent infinite loop + mock_paginate_aws_response.return_value = ( + [], # results + { # pagination_metadata + 'total_count': 0, + 'next_token': 'next-page-token', + 'has_more_results': True, + 'pages_fetched': 1, + }, + ) + + # Execute + result = await list_workload_estimates( + mock_context, next_token='current-token', max_pages=2 + ) + + # Assert + mock_paginate_aws_response.assert_called_once() + assert result['status'] == 'success' + assert result['data']['pagination']['next_token'] == 'next-page-token' + assert result['data']['pagination']['has_more_results'] is True + + @pytest.mark.asyncio + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.paginate_aws_response' + ) + async def test_list_workload_estimate_usage_with_pagination( + self, + mock_paginate_aws_response, + mock_create_client, + mock_get_preferences, + mock_context, + mock_bcm_pricing_calculator_client, + ): + """Test list_workload_estimate_usage with pagination token.""" + # Setup + mock_create_client.return_value = mock_bcm_pricing_calculator_client + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Mock the paginate_aws_response function to prevent infinite loop + mock_paginate_aws_response.return_value = ( + [], # results + { # pagination_metadata + 'total_count': 0, + 'next_token': 'usage-next-token', + 'has_more_results': True, + 'pages_fetched': 1, + }, + ) + + # Execute + result = await list_workload_estimate_usage( + mock_context, + workload_estimate_id='estimate-123', + next_token='usage-current-token', + max_pages=2, + ) + + # Assert + mock_paginate_aws_response.assert_called_once() + assert result['status'] == 'success' + assert result['data']['pagination']['next_token'] == 'usage-next-token' + assert result['data']['pagination']['has_more_results'] is True + + def test_format_workload_estimate_response_with_string_timestamps(self): + """Test format_workload_estimate_response with string timestamps.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'VALID', + 'createdAt': '2023-01-01T12:00:00Z', + 'expiresAt': '2023-12-31T23:59:59Z', + 'rateTimestamp': '2023-01-01T00:00:00Z', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert result['created_at']['timestamp'] == '2023-01-01T12:00:00Z' + assert result['created_at']['formatted'] == '2023-01-01T12:00:00Z' + assert result['expires_at']['timestamp'] == '2023-12-31T23:59:59Z' + assert result['rate_timestamp']['timestamp'] == '2023-01-01T00:00:00Z' + + def test_format_workload_estimate_response_none_cost(self): + """Test format_workload_estimate_response with None total cost.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': 'VALID', + 'totalCost': None, + 'costCurrency': 'USD', + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert 'cost' in result + assert result['cost']['amount'] is None + assert result['cost']['currency'] == 'USD' + assert result['cost']['formatted'] is None + + def test_format_usage_item_response_no_historical_bill_interval(self): + """Test format_usage_item_response with historical usage but no bill interval.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'historicalUsage': { + 'serviceCode': 'AmazonEC2', + 'usageType': 'BoxUsage:t3.medium', + 'operation': 'RunInstances', + 'location': 'US East (N. Virginia)', + 'usageAccountId': '123456789012', + # No billInterval + }, + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert 'historical_usage' in result + historical = result['historical_usage'] + assert 'bill_interval' not in historical + + def test_format_usage_item_response_empty_historical_usage(self): + """Test format_usage_item_response with empty historical usage.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'historicalUsage': {}, + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + # Empty historicalUsage dict should not add historical_usage to result + assert 'historical_usage' not in result + assert result['id'] == 'usage-123' + assert result['service_code'] == 'AmazonEC2' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.handle_aws_error' + ) + async def test_get_workload_estimate_exception( + self, mock_handle_error, mock_create_client, mock_get_preferences, mock_context + ): + """Test get_workload_estimate handles exceptions properly.""" + # Setup + error = ClientError( + {'Error': {'Code': 'ResourceNotFoundException', 'Message': 'Estimate not found'}}, + 'GetWorkloadEstimate', + ) + mock_get_preferences.return_value = {'account_types': ['management account']} + mock_create_client.return_value.get_workload_estimate.side_effect = error + mock_handle_error.return_value = {'status': 'error', 'message': 'Estimate not found'} + + # Execute + result = await get_workload_estimate(mock_context, identifier='nonexistent-estimate') + + # Assert + mock_handle_error.assert_called_once_with( + mock_context, error, 'get_workload_estimate', 'BCM Pricing Calculator' + ) + assert result['status'] == 'error' + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.create_aws_client' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.handle_aws_error' + ) + async def test_list_workload_estimate_usage_exception( + self, mock_handle_error, mock_create_client, mock_get_preferences, mock_context + ): + """Test list_workload_estimate_usage handles exceptions properly.""" + # Setup + error = BotoCoreError() + mock_get_preferences.return_value = {'account_types': ['management account']} + mock_create_client.return_value.list_workload_estimate_usage.side_effect = error + mock_handle_error.return_value = {'status': 'error', 'message': 'Connection error'} + + # Execute + result = await list_workload_estimate_usage( + mock_context, workload_estimate_id='estimate-123' + ) + + # Assert + mock_handle_error.assert_called_once_with( + mock_context, error, 'list_workload_estimate_usage', 'BCM Pricing Calculator' + ) + assert result['status'] == 'error' + + +@pytest.mark.parametrize( + 'status,expected_indicator', + [ + ('VALID', 'Valid'), + ('UPDATING', 'Updating'), + ('INVALID', 'Invalid'), + ('ACTION_NEEDED', 'Action Needed'), + ('UNKNOWN', '❓ UNKNOWN'), + ], +) +def test_format_workload_estimate_response_status_indicators(status, expected_indicator): + """Test format_workload_estimate_response status indicators.""" + # Setup + estimate = { + 'id': 'estimate-123', + 'name': 'Test Estimate', + 'status': status, + } + + # Execute + result = format_workload_estimate_response(estimate) + + # Assert + assert result['status_indicator'] == expected_indicator + + +@pytest.mark.parametrize( + 'status,expected_indicator', + [ + ('VALID', 'Valid'), + ('INVALID', 'Invalid'), + ('STALE', 'Stale'), + ('UNKNOWN', '❓ UNKNOWN'), + ], +) +def test_format_usage_item_response_status_indicators(status, expected_indicator): + """Test format_usage_item_response status indicators.""" + # Setup + usage_item = { + 'id': 'usage-123', + 'serviceCode': 'AmazonEC2', + 'status': status, + } + + # Execute + result = format_usage_item_response(usage_item) + + # Assert + assert result['status_indicator'] == expected_indicator + + +@pytest.mark.asyncio +class TestBcmPricingCalcCoreFunction: + """Tests for the core bcm_pricing_calc_core function (lines 101-149).""" + + async def test_bcm_pricing_calc_core_invalid_operation(self, mock_context): + """Test bcm_pricing_calc_core with invalid operation (covers lines 106-113).""" + # Execute - call the core function directly + result = await bcm_pricing_calc_core(mock_context, operation='invalid_operation') + + # Assert + assert result['status'] == 'error' + assert 'Invalid operation' in result['message'] + assert 'invalid_parameter' in result['data'] + mock_context.info.assert_called_with( + 'Received BCM Pricing Calculator operation: invalid_operation' + ) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_workload_estimate' + ) + async def test_bcm_pricing_calc_core_get_workload_estimate_operation( + self, mock_get_workload_estimate, mock_context + ): + """Test bcm_pricing_calc_core with get_workload_estimate operation (covers lines 115-117).""" + # Setup + mock_get_workload_estimate.return_value = { + 'status': 'success', + 'data': {'workload_estimate': {}}, + } + + # Execute - call the core function directly + result = await bcm_pricing_calc_core( + mock_context, operation='get_workload_estimate', identifier='estimate-123' + ) + + # Assert + mock_get_workload_estimate.assert_called_once_with(mock_context, 'estimate-123') + assert result['status'] == 'success' + mock_context.info.assert_called_with( + 'Received BCM Pricing Calculator operation: get_workload_estimate' + ) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.list_workload_estimates' + ) + async def test_bcm_pricing_calc_core_list_workload_estimates_operation( + self, mock_list_workload_estimates, mock_context + ): + """Test bcm_pricing_calc_core with list_workload_estimates operation (covers lines 118-121).""" + # Setup + mock_list_workload_estimates.return_value = { + 'status': 'success', + 'data': {'workload_estimates': []}, + } + + # Execute - call the core function directly + result = await bcm_pricing_calc_core( + mock_context, + operation='list_workload_estimates', + created_after='2023-01-01T00:00:00Z', + created_before='2023-12-31T23:59:59Z', + expires_after='2023-06-01T00:00:00Z', + expires_before='2024-01-01T00:00:00Z', + status_filter='VALID', + name_filter='Test', + name_match_option='CONTAINS', + next_token='token123', + max_results=50, + ) + + # Assert + mock_list_workload_estimates.assert_called_once_with( + mock_context, + '2023-01-01T00:00:00Z', + '2023-12-31T23:59:59Z', + '2023-06-01T00:00:00Z', + '2024-01-01T00:00:00Z', + 'VALID', + 'Test', + 'CONTAINS', + 'token123', + 50, + None, + ) + assert result['status'] == 'success' + mock_context.info.assert_called_with( + 'Received BCM Pricing Calculator operation: list_workload_estimates' + ) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.list_workload_estimate_usage' + ) + async def test_bcm_pricing_calc_core_list_workload_estimate_usage_operation( + self, mock_list_usage, mock_context + ): + """Test bcm_pricing_calc_core with list_workload_estimate_usage operation (covers lines 122-125).""" + # Setup + mock_list_usage.return_value = {'status': 'success', 'data': {'usage_items': []}} + + # Execute - call the core function directly + result = await bcm_pricing_calc_core( + mock_context, + operation='list_workload_estimate_usage', + identifier='estimate-123', + usage_account_id_filter='123456789012', + service_code_filter='AmazonEC2', + usage_type_filter='BoxUsage', + operation_filter='RunInstances', + location_filter='US East (N. Virginia)', + usage_group_filter='EC2-Instance', + next_token='usage-token', + max_results=100, + ) + + # Assert + mock_list_usage.assert_called_once_with( + mock_context, + 'estimate-123', + '123456789012', + 'AmazonEC2', + 'BoxUsage', + 'RunInstances', + 'US East (N. Virginia)', + 'EC2-Instance', + 'usage-token', + 100, + None, + ) + assert result['status'] == 'success' + mock_context.info.assert_called_with( + 'Received BCM Pricing Calculator operation: list_workload_estimate_usage' + ) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + async def test_bcm_pricing_calc_core_get_preferences_operation_success( + self, mock_get_preferences, mock_context + ): + """Test bcm_pricing_calc_core with get_preferences operation - success case (covers lines 126-133).""" + # Setup + mock_get_preferences.return_value = {'account_types': ['management account']} + + # Execute - call the core function directly + result = await bcm_pricing_calc_core(mock_context, operation='get_preferences') + + # Assert + mock_get_preferences.assert_called_once_with(mock_context) + assert result['status'] == 'success' + assert result['data']['message'] == 'Preferences are properly configured' + mock_context.info.assert_called_with( + 'Received BCM Pricing Calculator operation: get_preferences' + ) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_preferences' + ) + async def test_bcm_pricing_calc_core_get_preferences_operation_not_configured( + self, mock_get_preferences, mock_context + ): + """Test bcm_pricing_calc_core with get_preferences operation - not configured case (covers lines 127-131).""" + # Setup + mock_get_preferences.return_value = { + 'error': 'BCM Pricing Calculator preferences are not configured. Please configure preferences before using this service.' + } + + # Execute - call the core function directly + result = await bcm_pricing_calc_core(mock_context, operation='get_preferences') + + # Assert + mock_get_preferences.assert_called_once_with(mock_context) + assert result['status'] == 'error' + assert result['data']['error'] == PREFERENCES_NOT_CONFIGURED_ERROR + assert result['message'] == PREFERENCES_NOT_CONFIGURED_ERROR + mock_context.info.assert_called_with( + 'Received BCM Pricing Calculator operation: get_preferences' + ) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.get_workload_estimate' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.handle_aws_error' + ) + async def test_bcm_pricing_calc_core_exception_handling( + self, mock_handle_error, mock_get_workload_estimate, mock_context + ): + """Test bcm_pricing_calc_core exception handling (covers lines 137-149).""" + # Setup + test_error = Exception('Test error') + mock_get_workload_estimate.side_effect = test_error + mock_handle_error.return_value = { + 'data': {'error': 'Test error message'}, + 'status': 'error', + } + + # Execute - call the core function directly + result = await bcm_pricing_calc_core( + mock_context, operation='get_workload_estimate', identifier='estimate-123' + ) + + # Assert + mock_handle_error.assert_called_once_with( + mock_context, + test_error, + 'get_workload_estimate', + 'AWS Billing and Cost Management Pricing Calculator', + ) + mock_context.error.assert_called_once() + error_call_args = mock_context.error.call_args[0][0] + assert ( + 'Failed to process AWS Billing and Cost Management Pricing Calculator request' + in error_call_args + ) + assert 'Test error message' in error_call_args + + assert result['status'] == 'error' + assert result['data']['error'] == 'Test error message' + assert ( + 'Failed to process AWS Billing and Cost Management Pricing Calculator request' + in result['message'] + ) + + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.list_workload_estimates' + ) + @patch( + 'awslabs.billing_cost_management_mcp_server.tools.bcm_pricing_calculator_tools.handle_aws_error' + ) + async def test_bcm_pricing_calc_core_exception_handling_no_error_in_response( + self, mock_handle_error, mock_list_estimates, mock_context + ): + """Test bcm_pricing_calc_core exception handling when error response has no error field (covers lines 137-149).""" + # Setup + test_error = Exception('Direct error message') + mock_list_estimates.side_effect = test_error + mock_handle_error.return_value = { + 'data': {}, # No error field in data + 'status': 'error', + } + + # Execute - call the core function directly + result = await bcm_pricing_calc_core(mock_context, operation='list_workload_estimates') + + # Assert + mock_handle_error.assert_called_once_with( + mock_context, + test_error, + 'list_workload_estimates', + 'AWS Billing and Cost Management Pricing Calculator', + ) + mock_context.error.assert_called_once() + error_call_args = mock_context.error.call_args[0][0] + assert ( + 'Failed to process AWS Billing and Cost Management Pricing Calculator request' + in error_call_args + ) + assert 'Direct error message' in error_call_args + + assert result['status'] == 'error' + assert result['data']['error'] == 'Direct error message' + assert ( + 'Failed to process AWS Billing and Cost Management Pricing Calculator request' + in result['message'] + ) diff --git a/src/billing-cost-management-mcp-server/uv.lock b/src/billing-cost-management-mcp-server/uv.lock index 1220a36ba7..587d7116a3 100644 --- a/src/billing-cost-management-mcp-server/uv.lock +++ b/src/billing-cost-management-mcp-server/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10" [[package]] @@ -58,7 +58,7 @@ wheels = [ [[package]] name = "awslabs-billing-cost-management-mcp-server" -version = "0.0.0" +version = "0.0.2" source = { editable = "." } dependencies = [ { name = "boto3" },