mchades commented on code in PR #14208: URL: https://github.com/apache/gravitino/pull/14208#discussion_r4408002667
########## mcp-server/tests/integration/test_semantic_model_metric_discovery.py: ########## @@ -0,0 +1,343 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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. + +"""Discover metrics through the live MCP HTTP server and Gravitino REST API.""" + +import asyncio +from decimal import Decimal +from uuid import uuid4 + +import httpx +import pytest +import simplejson +from fastmcp import Client +from fastmcp.client.transports import StreamableHttpTransport +from fastmcp.exceptions import ToolError + +from tests.integration.gravitino_setup import basic_auth_header + +# Pytest injects fixtures by argument name. +# pylint: disable=redefined-outer-name + + +class MetricDiscoveryFixture: + """Isolated Semantic Models and role grants for each live-server test.""" + + def __init__(self, env, rest_client): + self.env = env + self.rest_client = rest_client + self.metalake = "metric_mcp_" + uuid4().hex + self.base = f"/api/metalakes/{self.metalake}" + self.models_path = ( + self.base + + "/catalogs/metrics_catalog/schemas/semantic/semantic-models" + ) + self.revenue = self.metric("Revenue") + self.lower_revenue = self.metric("revenue") + + def request(self, method, path, body=None): + response = self.rest_client.request( + method, + path, + content=( + None + if body is None + else simplejson.dumps(body, use_decimal=True) + ), + ) + response.raise_for_status() + return simplejson.loads(response.content, use_decimal=True) + + def provision(self): + self.request( + "POST", + self.base + "/catalogs", + { + "name": "metrics_catalog", + "type": "RELATIONAL", + "provider": "lakehouse-generic", + "comment": "", + "properties": {}, + }, + ) + schemas_path = self.base + "/catalogs/metrics_catalog/schemas" + for schema in ("source", "semantic", "other", "empty"): + self.request( + "POST", + schemas_path, + {"name": schema, "comment": "", "properties": {}}, + ) + self.request( + "POST", + schemas_path + "/source/tables", + { + "name": "orders", + "columns": [{"name": "amount", "type": "long"}], + "properties": { + "format": "delta", + "external": "true", + "location": f"file:///tmp/{self.metalake}/orders", + }, + }, + ) + for model, metrics in ( + ("sales", [self.lower_revenue, self.revenue]), + ("hidden", [self.revenue]), + ("no_metrics", []), + ): + self.request( + "POST", + self.models_path, + { + "name": model, + "definition": self.definition(metrics), + "properties": {}, + }, + ) + self.request( + "POST", + schemas_path + "/other/semantic-models", + { + "name": "sales", + "definition": self.definition([self.metric("other_schema")]), + }, + ) + self.request("POST", self.base + "/users", {"name": "metric_reader"}) + self.request( + "POST", + self.base + "/roles", + { + "name": "metric_reader", + "properties": {}, + "securableObjects": [ + { + "fullName": "metrics_catalog", + "type": "CATALOG", + "privileges": [ + {"name": "USE_CATALOG", "condition": "ALLOW"} + ], + }, + { + "fullName": "metrics_catalog.semantic", + "type": "SCHEMA", + "privileges": [ + {"name": "USE_SCHEMA", "condition": "ALLOW"} + ], + }, + ], + }, + ) + self.request( + "PUT", + self.base + "/permissions/users/metric_reader/grant", + {"roleNames": ["metric_reader"]}, + ) + + def call(self, tool, principal="admin", schema="semantic"): + async def invoke(): + transport = StreamableHttpTransport( + url=self.env["mcp_url"], + headers={"Authorization": basic_auth_header(principal)}, + ) + async with Client(transport) as client: + result = await client.call_tool( + tool, + { + "metalake": self.metalake, + "catalog_name": "metrics_catalog", + "schema_name": schema, + }, + ) + return simplejson.loads( + result.content[0].text, use_decimal=True + ) + + return asyncio.run(invoke()) + + def assert_discovery(self, expected, principal="admin"): + assert self.call("list_metric_infos", principal) == expected + assert self.call("list_metrics", principal) == [ Review Comment: These are synchronous pytest tests and fixtures, matching the existing authorization and OAuth integration tests. Each call creates and closes its MCP client within `asyncio.run()`, with no event loop already running in that test thread. Keeping the current pattern. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
