chore(mcp): extract shared chart helpers and ASCII rendering into separate modules (#39438)

tests/unit_tests/mcp_service/chart/test_chart_helpers.py

@@ -0,0 +1,108 @@
+# 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.
+
+from unittest.mock import MagicMock, patch
+
+from superset.mcp_service.chart.chart_helpers import (
+    extract_form_data_key_from_url,
+    find_chart_by_identifier,
+    get_cached_form_data,
+)
+
+
+def test_extract_form_data_key_from_url_with_key():
+    url = "http://localhost:8088/explore/?form_data_key=abc123&slice_id=1"
+    assert extract_form_data_key_from_url(url) == "abc123"
+
+
+def test_extract_form_data_key_from_url_no_key():
+    url = "http://localhost:8088/explore/?slice_id=1"
+    assert extract_form_data_key_from_url(url) is None
+
+
+def test_extract_form_data_key_from_url_none():
+    assert extract_form_data_key_from_url(None) is None
+
+
+def test_extract_form_data_key_from_url_empty():
+    assert extract_form_data_key_from_url("") is None
+
+
+def test_extract_form_data_key_from_url_multiple_params():
+    url = "http://localhost:8088/explore/?slice_id=5&form_data_key=xyz789&other=val"
+    assert extract_form_data_key_from_url(url) == "xyz789"
+
+
+@patch("superset.daos.chart.ChartDAO.find_by_id")
+def test_find_chart_by_identifier_int(mock_find):
+    mock_chart = MagicMock()
+    mock_chart.id = 42
+    mock_find.return_value = mock_chart
+
+    result = find_chart_by_identifier(42)
+    mock_find.assert_called_once_with(42)
+    assert result == mock_chart
+
+
+@patch("superset.daos.chart.ChartDAO.find_by_id")
+def test_find_chart_by_identifier_str_digit(mock_find):
+    mock_chart = MagicMock()
+    mock_find.return_value = mock_chart
+
+    result = find_chart_by_identifier("123")
+    mock_find.assert_called_once_with(123)
+    assert result == mock_chart
+
+
+@patch("superset.daos.chart.ChartDAO.find_by_id")
+def test_find_chart_by_identifier_uuid(mock_find):
+    mock_chart = MagicMock()
+    mock_find.return_value = mock_chart
+
+    uuid_str = "a1b2c3d4-5678-90ab-cdef-1234567890ab"
+    result = find_chart_by_identifier(uuid_str)
+    mock_find.assert_called_once_with(uuid_str, id_column="uuid")
+    assert result == mock_chart
+
+
+@patch("superset.daos.chart.ChartDAO.find_by_id")
+def test_find_chart_by_identifier_not_found(mock_find):
+    mock_find.return_value = None
+    result = find_chart_by_identifier(999)
+    assert result is None
+
+
+@patch(
+    "superset.commands.explore.form_data.get.GetFormDataCommand.run",
+    return_value='{"viz_type": "table"}',
+)
+@patch("superset.commands.explore.form_data.get.GetFormDataCommand.__init__")
+def test_get_cached_form_data_success(mock_init, mock_run):
+    mock_init.return_value = None
+    result = get_cached_form_data("test_key")
+    assert result == '{"viz_type": "table"}'
+
+
+@patch(
+    "superset.commands.explore.form_data.get.GetFormDataCommand.run",
+    side_effect=KeyError("not found"),
+)
+@patch("superset.commands.explore.form_data.get.GetFormDataCommand.__init__")
+def test_get_cached_form_data_key_error(mock_init, mock_run):
+    mock_init.return_value = None
+    result = get_cached_form_data("bad_key")
+    assert result is None

tests/unit_tests/mcp_service/chart/tool/test_update_chart.py

@@ -26,6 +26,7 @@ import pytest
 from fastmcp import Client
 
 from superset.mcp_service.app import mcp
+from superset.mcp_service.chart.chart_helpers import find_chart_by_identifier
 from superset.mcp_service.chart.chart_utils import DatasetValidationResult
 from superset.mcp_service.chart.schemas import (
     AxisConfig,
@@ -40,7 +41,6 @@ from superset.mcp_service.chart.schemas import (
 from superset.mcp_service.chart.tool.update_chart import (
     _build_preview_form_data,
     _build_update_payload,
-    _find_chart,
 )
 
 # The __init__.py re-exports the update_chart *function*, so a plain
@@ -519,7 +519,7 @@ class TestUpdateChartDatasetAccess:
 
 
 class TestFindChart:
-    """Tests for _find_chart helper."""
+    """Tests for find_chart_by_identifier helper (moved to chart_helpers)."""
 
     @patch("superset.daos.chart.ChartDAO.find_by_id")
     def test_find_chart_by_numeric_id(self, mock_find):
@@ -527,7 +527,7 @@ class TestFindChart:
         mock_chart = Mock()
         mock_find.return_value = mock_chart
 
-        result = _find_chart(42)
+        result = find_chart_by_identifier(42)
 
         mock_find.assert_called_once_with(42)
         assert result is mock_chart
@@ -538,7 +538,7 @@ class TestFindChart:
         mock_chart = Mock()
         mock_find.return_value = mock_chart
 
-        result = _find_chart("123")
+        result = find_chart_by_identifier("123")
 
         mock_find.assert_called_once_with(123)
         assert result is mock_chart
@@ -550,7 +550,7 @@ class TestFindChart:
         mock_find.return_value = mock_chart
 
         uuid = "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
-        result = _find_chart(uuid)
+        result = find_chart_by_identifier(uuid)
 
         mock_find.assert_called_once_with(uuid, id_column="uuid")
         assert result is mock_chart
@@ -560,7 +560,7 @@ class TestFindChart:
         """Returns None when chart is not found."""
         mock_find.return_value = None
 
-        result = _find_chart(999)
+        result = find_chart_by_identifier(999)
 
         assert result is None