|
1 | | -#!/usr/bin/env python3 |
2 | 1 | """MCP Server for generating OTel integration metric test files. |
3 | 2 |
|
4 | 3 | This server provides tools to generate test files similar to test_postgres_metrics.py |
5 | 4 | but for different integrations (Redis, MySQL, Kafka, etc.). |
6 | 5 | """ |
7 | 6 |
|
8 | 7 | import json |
| 8 | +import sys |
9 | 9 | from typing import Any |
10 | 10 |
|
11 | 11 | # MCP SDK imports |
|
14 | 14 | from mcp.types import Tool, TextContent, Resource |
15 | 15 | import mcp.server.stdio |
16 | 16 | except ImportError: |
17 | | - print("Error: MCP SDK not installed. Install with: pip install mcp") |
18 | | - exit(1) |
| 17 | + sys.exit(1) |
19 | 18 |
|
20 | 19 | # Path to reference test files |
21 | 20 | SYSTEM_TESTS_ROOT = Path(__file__).parent.parent.parent |
|
46 | 45 | "container_name": "mysql_container", |
47 | 46 | "smoke_test_operations": [ |
48 | 47 | "r = container.exec_run(\"mysql -u root -ppassword -e 'CREATE DATABASE IF NOT EXISTS test_db;'\")", |
49 | | - "r = container.exec_run(\"mysql -u root -ppassword test_db -e 'CREATE TABLE IF NOT EXISTS test_table (id INT PRIMARY KEY);'\")", |
| 48 | + 'r = container.exec_run("mysql -u root -ppassword test_db -e ' |
| 49 | + "'CREATE TABLE IF NOT EXISTS test_table (id INT PRIMARY KEY);'\")", |
50 | 50 | "r = container.exec_run(\"mysql -u root -ppassword test_db -e 'INSERT INTO test_table VALUES (1);'\")", |
51 | 51 | "logger.info(r.output)", |
52 | 52 | "r = container.exec_run(\"mysql -u root -ppassword test_db -e 'SELECT * FROM test_table;'\")", |
|
75 | 75 | "smoke_test_operations": [ |
76 | 76 | 'r = container.exec_run("kafka-topics --create --topic test-topic --bootstrap-server localhost:9092")', |
77 | 77 | "logger.info(r.output)", |
78 | | - 'r = container.exec_run("kafka-console-producer --topic test-topic --bootstrap-server localhost:9092", stdin="test message")', |
| 78 | + 'r = container.exec_run("kafka-console-producer --topic test-topic ' |
| 79 | + '--bootstrap-server localhost:9092", stdin="test message")', |
79 | 80 | ], |
80 | 81 | "expected_smoke_metrics": [ |
81 | 82 | "kafka.messages", |
@@ -131,7 +132,7 @@ def generate_test_file( |
131 | 132 | # Format expected smoke metrics |
132 | 133 | expected_metrics_formatted = ",\n ".join([f'"{m}"' for m in config["expected_smoke_metrics"]]) |
133 | 134 |
|
134 | | - template = f'''import time |
| 135 | + return f'''import time |
135 | 136 | from pathlib import Path |
136 | 137 | from typing import TYPE_CHECKING |
137 | 138 |
|
@@ -234,19 +235,15 @@ def test_main(self) -> None: |
234 | 235 | observed_metrics.add(metric) |
235 | 236 | logger.info(f" {{metric}} {{serie['points']}}") |
236 | 237 |
|
237 | | - all_metric_has_be_seen = True |
238 | 238 | for metric in expected_metrics: |
239 | 239 | if metric not in observed_metrics: |
240 | 240 | logger.error(f"Metric {{metric}} hasn't been observed") |
241 | 241 | all_metric_has_be_seen = False |
242 | 242 | else: |
243 | 243 | logger.info(f"Metric {{metric}} has been observed") |
244 | 244 |
|
245 | | - assert all_metric_has_be_seen |
246 | 245 | ''' |
247 | 246 |
|
248 | | - return template |
249 | | - |
250 | 247 |
|
251 | 248 | def generate_init_file() -> str: |
252 | 249 | """Generate __init__.py file.""" |
@@ -286,7 +283,10 @@ async def list_tools() -> list[Tool]: |
286 | 283 | }, |
287 | 284 | "feature_name": { |
288 | 285 | "type": "string", |
289 | | - "description": "Feature name for the @features decorator (optional, defaults to <integration>_receiver_metrics)", |
| 286 | + "description": ( |
| 287 | + "Feature name for the @features decorator " |
| 288 | + "(optional, defaults to <integration>_receiver_metrics)" |
| 289 | + ), |
290 | 290 | }, |
291 | 291 | }, |
292 | 292 | "required": ["integration_name", "metrics_json_file"], |
@@ -507,7 +507,7 @@ def test_main(self) -> None: |
507 | 507 |
|
508 | 508 |
|
509 | 509 | @app.call_tool() |
510 | | -async def call_tool(name: str, arguments: Any) -> list[TextContent]: |
| 510 | +async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: |
511 | 511 | """Handle tool calls.""" |
512 | 512 |
|
513 | 513 | if name == "generate_integration_test": |
@@ -537,7 +537,9 @@ async def call_tool(name: str, arguments: Any) -> list[TextContent]: |
537 | 537 | "shared_utility": { |
538 | 538 | "note": "Uses shared OtelMetricsValidator from utils/otel_metrics_validator.py", |
539 | 539 | "location": "utils/otel_metrics_validator.py", |
540 | | - "import_statement": "from utils.otel_metrics_validator import OtelMetricsValidator, get_collector_metrics_from_scenario", |
| 540 | + "import_statement": ( |
| 541 | + "from utils.otel_metrics_validator import OtelMetricsValidator, get_collector_metrics_from_scenario" |
| 542 | + ), |
541 | 543 | }, |
542 | 544 | "directory_structure": f""" |
543 | 545 | Create the following directory structure: |
@@ -600,7 +602,9 @@ async def call_tool(name: str, arguments: Any) -> list[TextContent]: |
600 | 602 | "shared_utility": { |
601 | 603 | "location": "utils/otel_metrics_validator.py", |
602 | 604 | "description": "Reusable metrics validation class for all OTel integration tests", |
603 | | - "import_statement": "from utils.otel_metrics_validator import OtelMetricsValidator, get_collector_metrics_from_scenario", |
| 605 | + "import_statement": ( |
| 606 | + "from utils.otel_metrics_validator import OtelMetricsValidator, get_collector_metrics_from_scenario" |
| 607 | + ), |
604 | 608 | }, |
605 | 609 | "classes": { |
606 | 610 | "OtelMetricsValidator": { |
@@ -650,7 +654,7 @@ async def call_tool(name: str, arguments: Any) -> list[TextContent]: |
650 | 654 | raise ValueError(f"Unknown tool: {name}") |
651 | 655 |
|
652 | 656 |
|
653 | | -async def main(): |
| 657 | +async def main() -> None: |
654 | 658 | """Main entry point for the MCP server.""" |
655 | 659 | async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): |
656 | 660 | await app.run( |
|
0 commit comments