bito-code-review[bot] commented on code in PR #43570:
URL: https://github.com/apache/superset/pull/43570#discussion_r4163535668
##########
tests/unit_tests/mcp_service/chart/test_chart_utils.py:
##########
@@ -160,6 +160,53 @@ def test_merge_bubble_preserves_omitted_defaults(
assert {key: merged[key] for key in expected} == expected
[email protected](
+ "updates,expected",
+ [
+ ({}, {"row_limit": 42}),
+ ({"row_limit": 200}, {"row_limit": 200}),
+ ],
+)
+def test_merge_heatmap_preserves_omitted_row_limit(
+ updates: dict[str, Any], expected: dict[str, Any]
+) -> None:
+ """Normalizing a heatmap config must not mark unsent fields as set.
+
+ ``merge_chart_form_data`` keeps an omitted row limit only when the field
+ is absent from ``model_fields_set``. Dumping the config without
+ ``exclude_unset`` hands back one where every field is set, so an update
+ that never mentions the row limit still resets it. (Heatmap has no
+ ``color_scheme`` control, so only the row limit applies here.)
+ """
+ from superset.mcp_service.chart.chart_utils import map_heatmap_config
+ from superset.mcp_service.chart.schemas import HeatmapChartConfig
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Inline duplicate imports</b></div>
<div id="fix">
`map_heatmap_config` and `HeatmapChartConfig` are already imported at module
level (lines 26, 47); these function-body imports duplicate them with no
circular-dependency justification or explanatory comment, contrary to the repo
rule requiring module-level imports. Extend the existing top-level import lists
instead.
</div>
</div>
<small><i>Code Review Run #e0a661</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
##########
tests/unit_tests/mcp_service/chart/test_chart_utils.py:
##########
@@ -160,6 +160,53 @@ def test_merge_bubble_preserves_omitted_defaults(
assert {key: merged[key] for key in expected} == expected
[email protected](
+ "updates,expected",
+ [
+ ({}, {"row_limit": 42}),
+ ({"row_limit": 200}, {"row_limit": 200}),
+ ],
+)
+def test_merge_heatmap_preserves_omitted_row_limit(
+ updates: dict[str, Any], expected: dict[str, Any]
+) -> None:
+ """Normalizing a heatmap config must not mark unsent fields as set.
+
+ ``merge_chart_form_data`` keeps an omitted row limit only when the field
+ is absent from ``model_fields_set``. Dumping the config without
+ ``exclude_unset`` hands back one where every field is set, so an update
+ that never mentions the row limit still resets it. (Heatmap has no
+ ``color_scheme`` control, so only the row limit applies here.)
+ """
+ from superset.mcp_service.chart.chart_utils import map_heatmap_config
+ from superset.mcp_service.chart.schemas import HeatmapChartConfig
+
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis=ColumnRef(name="product"),
+ y_axis=ColumnRef(name="revenue"),
+ metric=ColumnRef(name="revenue", aggregate="SUM"),
+ **updates,
+ )
+ config = DatasetValidator.normalize_column_names(
+ config,
+ dataset_id=1,
+ dataset_context=DatasetContext(
+ id=1,
+ table_name="sales",
+ database_name="db",
+ available_columns=[{"name": "Product"}, {"name": "Revenue"}],
+ available_metrics=[],
+ ),
+ )
+ new_form_data = map_heatmap_config(config)
+ existing = {"viz_type": new_form_data["viz_type"], "row_limit": 42}
+
+ merged = merge_chart_form_data(existing, new_form_data, config)
+
+ assert {key: merged[key] for key in expected} == expected
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Weak merge assertion</b></div>
<div id="fix">
The dict-comprehension assertion only inspects keys in `expected`, so a
`merge_chart_form_data` regression that dropped the update fields (`metric`,
`groupby`) would still pass. Sibling
`test_merge_chart_preserves_omitted_defaults` also asserts `merged["metric"] ==
new_form_data["metric"]`; mirror that here so the merge overlay contract is
actually validated.
</div>
</div>
<small><i>Code Review Run #e0a661</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
##########
tests/unit_tests/mcp_service/chart/test_heatmap_chart.py:
##########
@@ -0,0 +1,288 @@
+# 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.
+
+"""Tests for the heatmap chart type plugin.
+
+Schema validation, form_data mapping (matching the frontend Heatmap
+buildQuery contract for viz_type ``heatmap_v2`` — an ``x_axis`` column, a
+single ``groupby`` Y column, and one ``metric``), native ``groupby``
+aliasing for the Y axis, and registry integration.
+"""
+
+import pytest
+from pydantic import TypeAdapter, ValidationError
+
+from superset.mcp_service.chart.chart_utils import map_heatmap_config
+from superset.mcp_service.chart.schemas import ChartConfig, HeatmapChartConfig
+
+
+class TestHeatmapChartConfigSchema:
+ """HeatmapChartConfig schema validation."""
+
+ def test_basic_heatmap_config(self) -> None:
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+ assert config.x_axis.name == "day_of_week"
+ assert config.y_axis.name == "hour"
+ assert config.normalize_across == "heatmap" # frontend default
+
+ def test_heatmap_missing_x_axis(self) -> None:
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+
+ def test_heatmap_missing_y_axis(self) -> None:
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+
+ def test_heatmap_missing_metric(self) -> None:
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ )
+
+ def test_heatmap_rejects_extra_fields(self) -> None:
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ bogus=1,
+ )
+
+ def test_heatmap_axis_rejects_aggregate(self) -> None:
+ """An aggregate makes an axis metric-like; x_axis/y_axis are dims."""
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week", "aggregate": "COUNT"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour", "aggregate": "COUNT"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+
+ def test_heatmap_y_axis_rejects_saved_metric(self) -> None:
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "count", "saved_metric": True},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+
+ def test_heatmap_invalid_normalize_across_rejected(self) -> None:
+ with pytest.raises(ValidationError):
+ HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ normalize_across="diagonal",
+ )
+
+ def test_groupby_alias_for_y_axis(self) -> None:
+ """Superset-native 'groupby' is accepted for the Y-axis field."""
+ config = HeatmapChartConfig.model_validate(
+ {
+ "chart_type": "heatmap_v2",
+ "x_axis": {"name": "day_of_week"},
+ "groupby": {"name": "hour"},
+ "metric": {"name": "trips", "aggregate": "COUNT"},
+ }
+ )
+ assert config.y_axis.name == "hour"
+
+ def test_chart_config_union_dispatches_heatmap(self) -> None:
+ config = TypeAdapter(ChartConfig).validate_python(
+ {
+ "chart_type": "heatmap_v2",
+ "x_axis": {"name": "day_of_week"},
+ "y_axis": {"name": "hour"},
+ "metric": {"name": "trips", "aggregate": "COUNT"},
+ }
+ )
+ assert isinstance(config, HeatmapChartConfig)
+
+
+class TestMapHeatmapConfig:
+ """form_data mapping must match the frontend Heatmap buildQuery."""
+
+ def test_basic_heatmap_form_data(self) -> None:
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+ form_data = map_heatmap_config(config)
+ assert form_data["viz_type"] == "heatmap_v2"
+ assert form_data["x_axis"] == "day_of_week"
+ # Y axis uses the groupby key as a single column (control is
multi:false)
+ assert form_data["groupby"] == "hour"
+ assert form_data["metric"]["label"] == "COUNT(trips)"
+ assert form_data["normalize_across"] == "heatmap"
+
+ def test_heatmap_form_data_with_normalize_and_filters(self) -> None:
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ normalize_across="x",
+ filters=[{"column": "year", "op": "=", "value": 2026}],
+ )
+ form_data = map_heatmap_config(config)
+ assert form_data["normalize_across"] == "x"
+ assert form_data["adhoc_filters"], "filters must map to adhoc_filters"
+
+ def test_normalized_defaults_false(self) -> None:
+ # normalize_across has no visual effect on the frontend unless the
+ # 'normalized' flag is also set, so it must be threaded through.
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+ assert map_heatmap_config(config)["normalized"] is False
+
+ def test_normalized_true_maps_through(self) -> None:
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ normalize_across="x",
+ normalized=True,
+ )
+ assert map_heatmap_config(config)["normalized"] is True
+
+ def test_heatmap_saved_metric_maps_to_name_string(self) -> None:
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "avg_fare", "saved_metric": True},
+ )
+ assert map_heatmap_config(config)["metric"] == "avg_fare"
+
+
+class TestHeatmapQueryContext:
+ """The built query must GROUP BY both axes, not just the Y (groupby)
column.
+
+ map_heatmap_config emits X under 'x_axis' and Y under 'groupby'; the query
+ builder folds x_axis into the columns only for time-series viz types, so
+ heatmap_v2 needs an explicit fold or its X dimension is dropped.
+ """
+
+ def test_x_axis_reaches_group_by(self, monkeypatch) -> None:
+ from superset.mcp_service.chart import chart_helpers
+
+ monkeypatch.setattr(
+ chart_helpers,
+ "resolve_datasource_engine",
+ lambda datasource_id, datasource_type: "base",
+ )
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
+ )
+ form_data = map_heatmap_config(config)
+ queries = chart_helpers.build_query_dicts_from_form_data(form_data, 1,
"table")
+ columns = queries[0]["columns"]
+ assert "day_of_week" in columns, "x_axis must reach GROUP BY"
+ assert "hour" in columns, "y_axis must reach GROUP BY"
+
+ def test_x_axis_reaches_columns_from_form_data(self) -> None:
+ # The generate_chart compile check derives columns through a different
+ # builder (columns_from_form_data, used by the dashboard export path
+ # too), which must tolerate the scalar 'groupby' this mapper emits and
+ # carry both axes rather than crashing on str.copy().
+ from superset.common.form_data_query_context import
columns_from_form_data
+
+ config = HeatmapChartConfig(
+ chart_type="heatmap_v2",
+ x_axis={"name": "day_of_week"},
+ y_axis={"name": "hour"},
+ metric={"name": "trips", "aggregate": "COUNT"},
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Repeated config arrange block</b></div>
<div id="fix">
The four-line `HeatmapChartConfig(chart_type=..., x_axis=..., y_axis=...,
metric=...)` arrange block repeats ~18 times across the file.
`test_bubble_chart.py` centralizes the identical pattern in a `_base()` helper
(`BubbleChartConfig(**_base())`); a `_heatmap_base(**overrides)` helper would
localize future schema changes to one site.
</div>
</div>
<small><i>Code Review Run #e0a661</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]
