format changes

Author: bojbrook

mcp_servers/integration_test_generator/constants.py

@@ -1,5 +1,7 @@
+from pathlib import Path
+
 METRIC_TYPES = {"sum", "gauge"}
-GH_BASE_API = "https://api.github.com/repos/open-telemetry/opentelemetry-collector-contrib/"
+GH_BASE_API = "https://api.github.com/repos/open-telemetry/opentelemetry-collector-contrib"
 
 # Path to reference test files
 SYSTEM_TESTS_ROOT = Path(__file__).parent.parent.parent

mcp_servers/integration_test_generator/server.py

@@ -6,6 +6,7 @@
 """
 
 import json
+import sys
 from pathlib import Path
 from typing import Any
 import requests
@@ -14,15 +15,19 @@
 from templates.test_integration_file_template import get_test_file_template
 from templates.prompt_template import get_generate_with_reference_prompt
 from metric_operations_generator import generate_smoke_operations_from_metrics
+from utils._logger import logger
+import yaml
+import base64  # The content is base64-encoded per GitHub API
 
 # MCP SDK imports
 try:
     from mcp.server import Server
     from mcp.types import Tool, TextContent, Resource
+    from mcp.types import Prompt, PromptArgument
     import mcp.server.stdio
 except ImportError:
-    print("Error: MCP SDK not installed. Install with: pip install mcp")
-    exit(1)
+    logger.error("Error: MCP SDK not installed. Install with: pip install mcp")
+    sys.exit(1)
 
 
 def generate_test_file(
@@ -196,7 +201,7 @@ def generate_metrics_file(integration_name: str) -> str:
         "Accept": "application/vnd.github+json",
     }
 
-    response = requests.get(url, headers=headers)
+    response = requests.get(url, headers=headers, timeout=2)
     response.raise_for_status()
 
     releases = response.json()  # list[dict]
@@ -205,36 +210,32 @@ def generate_metrics_file(integration_name: str) -> str:
 
     latest = releases[0]
 
-    metaDataUrl = f"{constants.GH_BASE_API}/contents/receiver/{integration_name.lower()}receiver/metadata.yaml?ref={latest.get('tag_name')}"
+    receiver_path = f"receiver/{integration_name.lower()}receiver/metadata.yaml"
+    metadata_url = f"{constants.GH_BASE_API}/contents/{receiver_path}?ref={latest.get('tag_name')}"
 
-    response = requests.get(metaDataUrl, headers=headers)
+    response = requests.get(metadata_url, headers=headers, timeout=2)
     if response.status_code != 200:
         return f"Failed to fetch metadata.yaml for {integration_name}."
 
     metadata_content = response.json().get("content")
     if not metadata_content:
         return "No content found in metadata.yaml."
 
-    # The content is base64-encoded per GitHub API
-    import base64
-
     try:
         decoded_yaml = base64.b64decode(metadata_content).decode("utf-8")
     except Exception as e:
         return f"Error decoding metadata.yaml content: {e}"
 
     # Parse YAML to get metrics info
     try:
-        import yaml
-
         yaml_data = yaml.safe_load(decoded_yaml)
     except Exception as e:
         return f"Error parsing YAML: {e}"
 
     metric_template = {}
     metrics_dict = yaml_data.get("metrics", {})
     for metric_name, metric_info in metrics_dict.items():
-        metric_type = set(metric_info.keys()) & METRIC_TYPES
+        metric_type = set(metric_info.keys()) & constants.METRIC_TYPES
         metric_template[metric_name] = {
             "data_type": metric_type.pop() if metric_type else None,
             "description": metric_info.get("description", ""),
@@ -249,9 +250,8 @@ def generate_metrics_file(integration_name: str) -> str:
         return "There is already a metric file created. Please delete and try again"
     # Create the parent directory if it doesn't exist
     metric_file_path.parent.mkdir(parents=True, exist_ok=True)
-    f = open(metric_file_path, "x")
-    f.write(result_json)
-    f.close()
+    with open(metric_file_path, "x") as file:
+        file.write(result_json)
 
     return metric_template
 
@@ -299,7 +299,10 @@ async def list_tools() -> list[Tool]:
                     },
                     "feature_name": {
                         "type": "string",
-                        "description": "Feature name for the @features decorator (optional, defaults to <integration>_receiver_metrics)",
+                        "description": (
+                            "Feature name for the @features decorator "
+                            "(optional, defaults to <integration>_receiver_metrics)"
+                        ),
                     },
                 },
                 "required": ["integration_name", "metrics_json_file"],
@@ -351,7 +354,7 @@ async def list_resources() -> list[Resource]:
     if constants.MYSQL_TEST_PATH.exists():
         resources.append(
             Resource(
-                uri=f"file://{MYSQL_TEST_PATH}",
+                uri=f"file://{constants.MYSQL_TEST_PATH}",
                 name="MySQL Metrics Test (Reference)",
                 description="MySQL metrics test implementation following PostgreSQL patterns",
                 mimeType="text/x-python",
@@ -402,7 +405,6 @@ async def read_resource(uri: str) -> str:
 @app.list_prompts()
 async def list_prompts():
     """List available prompts."""
-    from mcp.types import Prompt, PromptArgument
 
     return [
         Prompt(