From 32ffe2c9a47715385d4af3ee6c6b9de6ddd8cb5c Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Thu, 21 May 2026 01:28:45 +0000 Subject: [PATCH] fix(mcp): fix three failing unit tests in annotation layer tools - AnnotationList.filters_applied used list[AnnotationFilter] which only allows col="short_descr", but ModelListCore injects a col="layer_id" ColumnOperator that fails Pydantic validation, causing KeyError on layer_id in the JSON response. Changed to list[ColumnOperator]. - test_get_annotation_layer_info_found: ModelGetInfoCore._find_object calls find_by_id(id, query_options=None), not find_by_id(id); updated assertion to assert_called_once_with(5, query_options=None). --- .../mcp_service/annotation_layer/schemas.py | 2 +- .../tool/test_annotation_layer_tools.py | 32 +++++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/superset/mcp_service/annotation_layer/schemas.py b/superset/mcp_service/annotation_layer/schemas.py index 2e44949daf5..175ee25aeec 100644 --- a/superset/mcp_service/annotation_layer/schemas.py +++ b/superset/mcp_service/annotation_layer/schemas.py @@ -194,7 +194,7 @@ class AnnotationList(BaseModel): columns_loaded: list[str] = Field(default_factory=list) columns_available: list[str] = Field(default_factory=list) sortable_columns: list[str] = Field(default_factory=list) - filters_applied: list[AnnotationFilter] = Field(default_factory=list) + filters_applied: list[ColumnOperator] = Field(default_factory=list) pagination: PaginationInfo | None = None timestamp: datetime | None = None model_config = ConfigDict(ser_json_timedelta="iso8601") diff --git a/tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py b/tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py index 2963accd349..7bef662bfb5 100644 --- a/tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py +++ b/tests/unit_tests/mcp_service/annotation_layer/tool/test_annotation_layer_tools.py @@ -74,7 +74,7 @@ def make_annotation( # --------------------------------------------------------------------------- -@pytest.fixture() +@pytest.fixture def mcp_server(): return mcp @@ -137,7 +137,7 @@ class TestAnnotationFilterSchema: @patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_list_annotation_layers_basic(mock_list, mcp_server): """Basic listing returns structured response with annotation layers.""" layer = make_layer() @@ -157,7 +157,7 @@ async def test_list_annotation_layers_basic(mock_list, mcp_server): @patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_list_annotation_layers_empty(mock_list, mcp_server): """Empty result set returns zero count.""" mock_list.return_value = ([], 0) @@ -171,7 +171,7 @@ async def test_list_annotation_layers_empty(mock_list, mcp_server): @patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_list_annotation_layers_search(mock_list, mcp_server): """Search parameter is passed through to DAO.""" layer = make_layer(name="Release Events") @@ -190,7 +190,7 @@ async def test_list_annotation_layers_search(mock_list, mcp_server): @patch("superset.daos.annotation_layer.AnnotationLayerDAO.list") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_list_annotation_layers_pagination(mock_list, mcp_server): """Pagination metadata is correctly computed.""" mock_list.return_value = ([], 50) @@ -217,7 +217,7 @@ async def test_list_annotation_layers_pagination(mock_list, mcp_server): @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") -@pytest.mark.asyncio() +@pytest.mark.asyncio 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") @@ -231,11 +231,11 @@ async def test_get_annotation_layer_info_found(mock_find, mcp_server): data = json.loads(result.content[0].text) assert data["id"] == 5 assert data["name"] == "Prod Events" - mock_find.assert_called_once_with(5) + mock_find.assert_called_once_with(5, query_options=None) @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") -@pytest.mark.asyncio() +@pytest.mark.asyncio 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 @@ -258,7 +258,7 @@ async def test_get_annotation_layer_info_not_found(mock_find, mcp_server): @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") @patch("superset.daos.annotation_layer.AnnotationDAO.list") -@pytest.mark.asyncio() +@pytest.mark.asyncio 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) @@ -280,7 +280,7 @@ async def test_list_layer_annotations_basic(mock_list, mock_layer_find, mcp_serv @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") @patch("superset.daos.annotation_layer.AnnotationDAO.list") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_list_layer_annotations_layer_id_filter_prepended( mock_list, mock_layer_find, mcp_server ): @@ -308,7 +308,7 @@ async def test_list_layer_annotations_layer_id_filter_prepended( @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") -@pytest.mark.asyncio() +@pytest.mark.asyncio 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 @@ -326,7 +326,7 @@ async def test_list_layer_annotations_layer_not_found(mock_layer_find, mcp_serve @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") @patch("superset.daos.annotation_layer.AnnotationDAO.list") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_list_layer_annotations_only_returns_own_layer( mock_list, mock_layer_find, mcp_server ): @@ -354,7 +354,7 @@ async def test_list_layer_annotations_only_returns_own_layer( @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") @patch("superset.daos.annotation_layer.AnnotationDAO.find_by_id") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_get_layer_annotation_info_found( mock_ann_find, mock_layer_find, mcp_server ): @@ -375,7 +375,7 @@ async def test_get_layer_annotation_info_found( @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_get_layer_annotation_info_layer_not_found(mock_layer_find, mcp_server): """Returns error when the layer does not exist.""" mock_layer_find.return_value = None @@ -393,7 +393,7 @@ async def test_get_layer_annotation_info_layer_not_found(mock_layer_find, mcp_se @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") @patch("superset.daos.annotation_layer.AnnotationDAO.find_by_id") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_get_layer_annotation_info_annotation_not_found( mock_ann_find, mock_layer_find, mcp_server ): @@ -414,7 +414,7 @@ async def test_get_layer_annotation_info_annotation_not_found( @patch("superset.daos.annotation_layer.AnnotationLayerDAO.find_by_id") @patch("superset.daos.annotation_layer.AnnotationDAO.find_by_id") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_get_layer_annotation_info_wrong_layer( mock_ann_find, mock_layer_find, mcp_server ):