Files
superset2/tests/unit_tests/mcp_service/utils/test_response_utils.py

110 lines
4.2 KiB
Python

# 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.
"""
Unit tests for MCP service response utilities.
"""
from typing import Any
from superset.mcp_service.utils.response_utils import (
format_data_columns,
STATS_ROW_CAP,
)
class TestFormatDataColumns:
"""Test format_data_columns function."""
def test_infers_numeric_type(self) -> None:
"""Should infer numeric data_type from sample values."""
data: list[dict[str, Any]] = [
{"revenue": 100},
{"revenue": 200},
{"revenue": None},
]
columns = format_data_columns(data, ["revenue"])
assert len(columns) == 1
assert columns[0].name == "revenue"
assert columns[0].data_type == "numeric"
assert columns[0].null_count == 1
assert columns[0].unique_count == 2
def test_infers_boolean_type(self) -> None:
"""Should infer boolean data_type from sample values."""
data: list[dict[str, Any]] = [{"is_active": True}, {"is_active": False}]
columns = format_data_columns(data, ["is_active"])
assert columns[0].data_type == "boolean"
def test_infers_string_type_by_default(self) -> None:
"""Should default to string data_type when values aren't numeric/boolean."""
data: list[dict[str, Any]] = [{"region": "west"}, {"region": "east"}]
columns = format_data_columns(data, ["region"])
assert columns[0].data_type == "string"
def test_string_type_when_no_sample_values(self) -> None:
"""Should default to string data_type when all sampled values are null."""
data: list[dict[str, Any]] = [{"region": None}]
columns = format_data_columns(data, ["region"])
assert columns[0].data_type == "string"
def test_sample_values_capped_at_three(self) -> None:
"""Should only take the first 3 non-null values as samples."""
data: list[dict[str, Any]] = [{"region": f"r{i}"} for i in range(10)]
columns = format_data_columns(data, ["region"])
assert columns[0].sample_values == ["r0", "r1", "r2"]
def test_null_and_unique_counts_reflect_full_small_dataset(self) -> None:
"""Should count nulls/uniques across all rows when under the cap."""
data: list[dict[str, Any]] = [
{"region": "west"},
{"region": "west"},
{"region": "east"},
{"region": None},
]
columns = format_data_columns(data, ["region"])
assert columns[0].null_count == 1
assert columns[0].unique_count == 2
assert columns[0].statistics is None
def test_stats_marked_as_sampled_beyond_row_cap(self) -> None:
"""Should mark statistics as sampled when data exceeds STATS_ROW_CAP.
Regression test: null_count/unique_count are computed on a capped
sample for performance, but were previously returned as if they were
exact full-dataset totals with no indication of sampling.
"""
data: list[dict[str, Any]] = [{"id": i} for i in range(STATS_ROW_CAP + 10)]
columns = format_data_columns(data, ["id"])
assert columns[0].statistics == {"sampled_rows": STATS_ROW_CAP}
assert columns[0].null_count == 0
assert columns[0].unique_count == STATS_ROW_CAP
def test_multiple_columns(self) -> None:
"""Should build metadata for every requested column."""
data = [{"region": "west", "revenue": 100}]
columns = format_data_columns(data, ["region", "revenue"])
assert [c.name for c in columns] == ["region", "revenue"]