bito-code-review[bot] commented on code in PR #40342: URL: https://github.com/apache/superset/pull/40342#discussion_r3292313001
########## tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py: ########## @@ -0,0 +1,434 @@ +# 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. + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp import Client +from pydantic import ValidationError + +from superset.mcp_service.annotation_layer.schemas import ( + AnnotationFilter, + AnnotationLayerFilter, + ListAnnotationLayersRequest, + ListLayerAnnotationsRequest, +) +from superset.mcp_service.app import mcp +from superset.utils import json + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def make_layer( + layer_id: int = 1, name: str = "My Layer", descr: str = "desc" +) -> MagicMock: + obj = MagicMock() + obj.id = layer_id + obj.name = name + obj.descr = descr + obj.changed_on = None + obj.created_on = None + return obj + + +def make_annotation( + annotation_id: int = 10, + layer_id: int = 1, + short_descr: str = "Deploy", + long_descr: str = "Deployment annotation", +) -> MagicMock: + obj = MagicMock() + obj.id = annotation_id + obj.layer_id = layer_id + obj.short_descr = short_descr + obj.long_descr = long_descr + obj.start_dttm = None + obj.end_dttm = None + obj.json_metadata = None + return obj + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + [email protected] +def mcp_server(): + return mcp + + [email protected](autouse=True) +def mock_auth(): + from unittest.mock import Mock + + with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user: + mock_user = Mock() + mock_user.id = 1 + mock_user.username = "admin" + mock_get_user.return_value = mock_user + yield mock_get_user + + +# --------------------------------------------------------------------------- +# Schema validation tests +# --------------------------------------------------------------------------- + + +class TestAnnotationLayerFilterSchema: + def test_valid_name_filter(self): + f = AnnotationLayerFilter(col="name", opr="eq", value="My Layer") + assert f.col == "name" + + def test_invalid_column_rejected(self): + with pytest.raises(ValidationError): + AnnotationLayerFilter(col="descr", opr="eq", value="x") + + def test_search_and_filters_mutual_exclusion(self): + with pytest.raises(ValidationError): + ListAnnotationLayersRequest( + search="foo", + filters=[{"col": "name", "opr": "eq", "value": "bar"}], + ) + + +class TestAnnotationFilterSchema: + def test_valid_short_descr_filter(self): + f = AnnotationFilter(col="short_descr", opr="eq", value="Deploy") + assert f.col == "short_descr" + + def test_invalid_column_rejected(self): + with pytest.raises(ValidationError): + AnnotationFilter(col="layer_id", opr="eq", value=1) + + def test_search_and_filters_mutual_exclusion(self): + with pytest.raises(ValidationError): + ListLayerAnnotationsRequest( + layer_id=1, + search="foo", + filters=[{"col": "short_descr", "opr": "eq", "value": "bar"}], + ) + + +# --------------------------------------------------------------------------- +# list_annotation_layers tests +# --------------------------------------------------------------------------- + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_basic(mock_list, mcp_server): + """Basic listing returns structured response with annotation layers.""" + layer = make_layer() + mock_list.return_value = ([layer], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_annotation_layers", + {"request": {"page": 1, "page_size": 10}}, + ) + + data = json.loads(result.content[0].text) + assert data["annotation_layers"] is not None + assert len(data["annotation_layers"]) == 1 + assert data["annotation_layers"][0]["id"] == 1 + assert data["annotation_layers"][0]["name"] == "My Layer" + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_empty(mock_list, mcp_server): + """Empty result set returns zero count.""" + mock_list.return_value = ([], 0) + + async with Client(mcp_server) as client: + result = await client.call_tool("list_annotation_layers", {}) + + data = json.loads(result.content[0].text) + assert data["annotation_layers"] == [] + assert data["total_count"] == 0 + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_search(mock_list, mcp_server): + """Search parameter is passed through to DAO.""" + layer = make_layer(name="Release Events") + mock_list.return_value = ([layer], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_annotation_layers", + {"request": {"search": "release"}}, + ) + + data = json.loads(result.content[0].text) + assert data["annotation_layers"][0]["name"] == "Release Events" + call_kwargs = mock_list.call_args.kwargs + assert call_kwargs["search"] == "release" + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") [email protected] +async def test_list_annotation_layers_pagination(mock_list, mcp_server): + """Pagination metadata is correctly computed.""" + mock_list.return_value = ([], 50) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_annotation_layers", + {"request": {"page": 2, "page_size": 25}}, + ) + + data = json.loads(result.content[0].text) + assert data["page"] == 2 + assert data["page_size"] == 25 + assert data["total_count"] == 50 + assert data["total_pages"] == 2 + # Page 2 of 2, so no next page + assert data["has_next"] is False + assert data["has_previous"] is True + + +# --------------------------------------------------------------------------- +# get_annotation_layer_info tests +# --------------------------------------------------------------------------- + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") [email protected] +async def test_get_annotation_layer_info_found(mock_find, mcp_server): + """Returns annotation layer data when found.""" + mock_find.return_value = make_layer(layer_id=5, name="Prod Events") + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_annotation_layer_info", + {"request": {"id": 5}}, + ) + + data = json.loads(result.content[0].text) + assert data["id"] == 5 + assert data["name"] == "Prod Events" + mock_find.assert_called_once_with(5, query_options=None) + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") [email protected] +async def test_get_annotation_layer_info_not_found(mock_find, mcp_server): + """Returns error response when layer is not found.""" + mock_find.return_value = None + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_annotation_layer_info", + {"request": {"id": 999}}, + ) + + data = json.loads(result.content[0].text) + assert data["error_type"] == "not_found" + assert "999" in data["error"] + + +# --------------------------------------------------------------------------- +# list_layer_annotations tests +# --------------------------------------------------------------------------- + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") +@patch("superset.daos.annotation_layer.AnnotationDAO.list") [email protected] +async def test_list_layer_annotations_basic(mock_list, mock_layer_find, mcp_server): + """Annotations are listed and scoped to the specified layer.""" + mock_layer_find.return_value = make_layer(layer_id=1) + ann = make_annotation(annotation_id=10, layer_id=1) + mock_list.return_value = ([ann], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 1, "page": 1, "page_size": 10}}, + ) + + data = json.loads(result.content[0].text) + assert data["layer_id"] == 1 + assert len(data["annotations"]) == 1 + assert data["annotations"][0]["id"] == 10 + assert data["annotations"][0]["layer_id"] == 1 + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") +@patch("superset.daos.annotation_layer.AnnotationDAO.list") [email protected] +async def test_list_layer_annotations_layer_id_filter_prepended( + mock_list, mock_layer_find, mcp_server +): + """The layer_id filter is always prepended to DAO column_operators.""" + mock_layer_find.return_value = make_layer(layer_id=3) + mock_list.return_value = ([], 0) + + async with Client(mcp_server) as client: + await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 3}}, + ) + + call_kwargs = mock_list.call_args.kwargs + filters = call_kwargs.get("column_operators", []) + # First filter must be the layer_id eq filter + assert filters, "Expected at least one filter (layer_id)" + first = filters[0] + col = first.get("col") if isinstance(first, dict) else getattr(first, "col", None) + val = ( + first.get("value") if isinstance(first, dict) else getattr(first, "value", None) + ) + assert col == "layer_id" + assert val == 3 + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") [email protected] +async def test_list_layer_annotations_layer_not_found(mock_layer_find, mcp_server): + """Returns error when the layer does not exist.""" + mock_layer_find.return_value = None + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 42}}, + ) + + data = json.loads(result.content[0].text) + assert data["error_type"] == "not_found" + assert "42" in data["error"] + + +@patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") +@patch("superset.daos.annotation_layer.AnnotationDAO.list") [email protected] +async def test_list_layer_annotations_only_returns_own_layer( + mock_list, mock_layer_find, mcp_server +): + """Results are filtered to the requested layer only — wrong layer_id is rejected.""" + mock_layer_find.return_value = make_layer(layer_id=1) + # Simulate DAO returning annotations — the layer_id filter is applied at DB level + ann_wrong = make_annotation(annotation_id=99, layer_id=2) + mock_list.return_value = ([ann_wrong], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_layer_annotations", + {"request": {"layer_id": 1}}, + ) + + data = json.loads(result.content[0].text) + # layer_id in response header must still be 1 (the requested layer) + assert data["layer_id"] == 1 Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Weak filter behavior assertion</b></div> <div id="fix"> The test `test_list_layer_annotations_only_returns_own_layer` asserts `data['layer_id'] == 1` (the request header), but never verifies that returned annotations belong to that layer. If the DAO-level filter fails silently, wrong annotations would pass this test. Add assertion like `assert data['annotations'][0]['layer_id'] == 1` to verify actual filtering. </div> <details> <summary> <b>Code suggestion</b> </summary> <blockquote>Check the AI-generated fix before applying</blockquote> <div id="code"> ``` --- tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py +++ tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py @@ -343,5 +343,7 @@ class TestListLayerAnnotations: data = json.loads(result.content[0].text) # layer_id in response header must still be 1 (the requested layer) assert data["layer_id"] == 1 + # Verify the returned annotation actually belongs to the requested layer + assert data["annotations"][0]["layer_id"] == 1 ``` </div> </details> </div> <small><i>Code Review Run #cc39a6</i></small> </div> --- Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>) - [ ] Yes, avoid them -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
