MCPX MCP Gateway

# type: ignore import json from datetime import datetime from typing import Any from aiohttp import ClientSession from prometheus_client.parser import text_string_to_metric_families from prometheus_client import Metric from behave import then, when from behave.api.async_step import async_run_until_complete from toolkit_testing.integration_tests.routing import Routing from toolkit_testing.integration_tests.s3 import AWSAccess, S3ClientHelper from utils.consts import * from utils.policies import EndpointPolicy, PoliciesRequests _minio_client_helper = S3ClientHelper( AWSAccess(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), Routing("127.0.0.1", 9000), ) _LUNAR_TRANSACTION_HISTOGRAM_METRIC_NAME = "lunar_transaction" _HISTOGRAM_COUNT_SUFFIX = "_count" _COUNTER_SUFFIX = "_total" @when( "policies.yaml includes a metrics_collector diagnosis for {method} {host} {path:Path} requests with {exporter} as exporter" ) def step_impl(context: Any, method: str, host: str, path: str, exporter: str): policies_requests: PoliciesRequests = context.policies_requests diagnosis = { "name": "metrics_collector for {method} {host}/{path}", "enabled": True, "export": exporter, "config": {"metrics_collector": {}}, } policies_requests.endpoints.append( EndpointPolicy(method, f"{host}{path}", diagnosis=[diagnosis]) ) @when( "policies.yaml includes a global metrics_collector diagnosis with {exporter} as exporter" ) def step_impl(context: Any, exporter: str): policies_requests: PoliciesRequests = context.policies_requests diagnosis = { "name": "global metrics_collector", "enabled": True, "export": exporter, "config": {"metrics_collector": {}}, } policies_requests.globals.diagnosis.append(diagnosis) @when( "policies.yaml includes a global metrics_collector diagnosis with {exporter} as exporter and custom counter for {response_header} response header" ) def step_impl(context: Any, exporter: str, response_header: str): policies_requests: PoliciesRequests = context.policies_requests diagnosis = { "name": "global metrics_collector", "enabled": True, "export": exporter, "config": { "metrics_collector": { "counters": [ { "name_suffix": response_header, "payload": "response_headers", "key": response_header, } ] } }, } policies_requests.globals.diagnosis.append(diagnosis) @then("Transaction metrics are written") @async_run_until_complete async def step_impl(context: Any): content_bytes = _minio_client_helper.get_object_content_of_last_modified_object( LUNAR_BUCKET_NAME, retries=20 ) assert content_bytes is not None content = content_bytes.decode() collected_metrics = _read_last_collected_metrics(content) print("content is: " + content) context.collected_metrics = collected_metrics @then("Traffic Metrics status should be {status_code:Int}") @async_run_until_complete async def step_impl(context: Any, status_code: int): assert context.collected_metrics["status_code"] == status_code @then("Traffic Metrics normalized_url should be {url}") @async_run_until_complete async def step_impl(context: Any, url: str): assert context.collected_metrics["normalized_url"] == url @then("Traffic Metrics method should be {method}") @async_run_until_complete async def step_impl(context: Any, method: str): assert context.collected_metrics["method"] == method @then("There are {count:Int} lunar_transaction histograms on Prometheus Metric Server") @async_run_until_complete async def step_impl(context: Any, count: int): url = f"http://localhost:{PROMETHEUS_METRIC_SERVER_PORT}{PROMETHEUS_METRICS_ROUTE}" async with ClientSession() as session: try: async with session.get(url=url) as resp: assert resp.status == 200 raw_metrics = await resp.text() except Exception as ex: print(f"failed calling metrics server: {ex}") assert False print("raw metrics:") print(raw_metrics) print("***") metrics = text_string_to_metric_families(raw_metrics) matched_metric = "" for metric in list(metrics): if metric.name == _LUNAR_TRANSACTION_HISTOGRAM_METRIC_NAME: matched_metric = metric break counts = [ sample for sample in matched_metric.samples if sample.name.endswith(_HISTOGRAM_COUNT_SUFFIX) ] context.histogram_metric = matched_metric assert len(counts) == count @then( "There is a histogram of status {code}, normalized_url {normalized_url} with {expected_count:Int} calls" ) @async_run_until_complete async def step_impl(context: Any, code: str, normalized_url: str, expected_count: int): histogram_metric: Metric = context.histogram_metric for sample in histogram_metric.samples: if ( sample.name.endswith(_HISTOGRAM_COUNT_SUFFIX) and sample.labels.get("status_code") == code and sample.labels.get("normalized_url") == normalized_url ): matched_sample = sample break assert matched_sample.value == float(expected_count) @then("There is a counter named {name} with the value {expected_value:Int}") @async_run_until_complete async def step_impl(context: Any, name: str, expected_value: int): url = f"http://localhost:{PROMETHEUS_METRIC_SERVER_PORT}{PROMETHEUS_METRICS_ROUTE}" async with ClientSession() as session: try: async with session.get(url=url) as resp: assert resp.status == 200 raw_metrics = await resp.text() except Exception as ex: print(f"failed calling metrics server: {ex}") assert False print("raw metrics:") print(raw_metrics) print("***") print("Parsing metrics...") metrics = text_string_to_metric_families(raw_metrics) print("Successfully parsed metrics") for metric in list(metrics): if metric.name == name.rstrip(_COUNTER_SUFFIX) or metric.name == name: matched_metric = metric break print("***") print(matched_metric) print("***") counter: Metric = matched_metric.samples[0] if counter.value != float(expected_value): print(f"expected value: {expected_value}, actual value: {counter.value}") assert counter.value == float(expected_value) @then("There is a gauge {name} with the value larger than {expected_value:Int}") @async_run_until_complete async def step_impl(context: Any, name: str, expected_value: int): url = f"http://localhost:{PROMETHEUS_METRIC_SERVER_PORT}{PROMETHEUS_METRICS_ROUTE}" async with ClientSession() as session: try: async with session.get(url=url) as resp: assert resp.status == 200 raw_metrics = await resp.text() except Exception as ex: print(f"failed calling metrics server: {ex}") assert False print("raw metrics:") print(raw_metrics) print("***") print("Parsing metrics...") metrics = text_string_to_metric_families(raw_metrics) print("Successfully parsed metrics") for metric in list(metrics): if metric.name == name.rstrip(_COUNTER_SUFFIX) or metric.name == name: matched_metric = metric break print("***") print(matched_metric) print("***") counter: Metric = matched_metric.samples[0] assert counter.value > float(expected_value) def _read_last_collected_metrics(content: str | None) -> dict[str, Any]: if content is None: return {} collected_metrics_data = content.splitlines() assert len(collected_metrics_data) > 0 raw_collected_metrics = collected_metrics_data[-1] return json.loads(raw_collected_metrics)