aminghadersohi commented on code in PR #41924: URL: https://github.com/apache/superset/pull/41924#discussion_r3573735141
########## tests/unit_tests/mcp_service/utils/test_retry_utils.py: ########## @@ -0,0 +1,396 @@ +# 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 retry utilities. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from sqlalchemy.exc import OperationalError + +from superset.mcp_service.utils.retry_utils import ( + async_retry_database_operation, + async_retry_on_exception, + exponential_backoff, + retry_database_operation, + retry_on_exception, + retry_screenshot_operation, + RetryableOperation, +) + +SLEEP = "superset.mcp_service.utils.retry_utils.time.sleep" Review Comment: Added the explicit `str` annotation to `SLEEP` in 8ab6108ae9. ########## tests/unit_tests/mcp_service/utils/test_retry_utils.py: ########## @@ -0,0 +1,396 @@ +# 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 retry utilities. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from sqlalchemy.exc import OperationalError + +from superset.mcp_service.utils.retry_utils import ( + async_retry_database_operation, + async_retry_on_exception, + exponential_backoff, + retry_database_operation, + retry_on_exception, + retry_screenshot_operation, + RetryableOperation, +) + +SLEEP = "superset.mcp_service.utils.retry_utils.time.sleep" +ASYNC_SLEEP = "superset.mcp_service.utils.retry_utils.asyncio.sleep" Review Comment: Added the explicit `str` annotation to `ASYNC_SLEEP` in 8ab6108ae9. ########## tests/unit_tests/mcp_service/utils/test_retry_utils.py: ########## @@ -0,0 +1,396 @@ +# 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 retry utilities. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from sqlalchemy.exc import OperationalError + +from superset.mcp_service.utils.retry_utils import ( + async_retry_database_operation, + async_retry_on_exception, + exponential_backoff, + retry_database_operation, + retry_on_exception, + retry_screenshot_operation, + RetryableOperation, +) + +SLEEP = "superset.mcp_service.utils.retry_utils.time.sleep" +ASYNC_SLEEP = "superset.mcp_service.utils.retry_utils.asyncio.sleep" + + +# --------------------------------------------------------------------------- +# exponential_backoff +# --------------------------------------------------------------------------- + + +def test_exponential_backoff_without_jitter_doubles_each_attempt() -> None: + """Should double the delay for each subsequent attempt when jitter is off.""" + assert exponential_backoff(0, base_delay=1.0, jitter=False) == 1.0 + assert exponential_backoff(1, base_delay=1.0, jitter=False) == 2.0 + assert exponential_backoff(2, base_delay=1.0, jitter=False) == 4.0 + + +def test_exponential_backoff_caps_at_max_delay() -> None: + """Should never return a delay larger than max_delay.""" + delay = exponential_backoff(10, base_delay=1.0, max_delay=5.0, jitter=False) + assert delay == 5.0 + + +def test_exponential_backoff_jitter_stays_within_twenty_five_percent() -> None: + """Should keep jittered delays within +/-25% of the base delay.""" + for _ in range(50): + delay = exponential_backoff(2, base_delay=1.0, jitter=True) + # Unjittered delay for attempt=2, base_delay=1.0 is 4.0. + assert 3.0 <= delay <= 5.0 + + +def test_exponential_backoff_never_negative() -> None: + """Should clamp the delay to be non-negative even with jitter applied.""" + for _ in range(50): + delay = exponential_backoff(0, base_delay=0.001, jitter=True) + assert delay >= 0 + + +# --------------------------------------------------------------------------- +# retry_on_exception (sync) +# --------------------------------------------------------------------------- + + +def test_retry_on_exception_succeeds_first_try() -> None: + """Should call the wrapped function once when it succeeds immediately.""" + mock_func = MagicMock(return_value="ok") + wrapped = retry_on_exception(max_attempts=3)(mock_func) + + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 1 + + +def test_retry_on_exception_retries_then_succeeds() -> None: + """Should retry on retryable exceptions and return the eventual success.""" + mock_func = MagicMock( + side_effect=[ConnectionError("fail"), ConnectionError("fail"), "ok"] + ) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, base_delay=0.01, jitter=False)( + mock_func + ) + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 3 + assert mock_sleep.call_count == 2 + + +def test_retry_on_exception_exhausts_retries_and_raises_last_exception() -> None: + """Should raise the last retryable exception once max_attempts is reached.""" + mock_func = MagicMock(side_effect=ConnectionError("always fails")) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, base_delay=0.01, jitter=False)( + mock_func + ) + with pytest.raises(ConnectionError, match="always fails"): + wrapped() + + assert mock_func.call_count == 3 + # No sleep after the final (failed) attempt. + assert mock_sleep.call_count == 2 + + +def test_retry_on_exception_non_retryable_exception_fails_immediately() -> None: + """Should not retry exceptions outside the configured retryable tuple.""" + mock_func = MagicMock(side_effect=ValueError("bad input")) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, exceptions=(ConnectionError,))( + mock_func + ) + with pytest.raises(ValueError, match="bad input"): + wrapped() + + assert mock_func.call_count == 1 + mock_sleep.assert_not_called() + + +def test_retry_on_exception_respects_custom_exceptions_tuple() -> None: + """Should retry only on the exception types passed via `exceptions`.""" + mock_func = MagicMock(side_effect=[TimeoutError("timed out"), "ok"]) + mock_func.__name__ = "mock_func" + + with patch(SLEEP): + wrapped = retry_on_exception( + max_attempts=2, base_delay=0.01, jitter=False, exceptions=(TimeoutError,) + )(mock_func) + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 2 + + +def test_retry_on_exception_max_attempts_zero_raises_runtime_error() -> None: + """Should raise RuntimeError when max_attempts leaves no attempt to run.""" + mock_func = MagicMock(return_value="ok") + mock_func.__name__ = "mock_func" + wrapped = retry_on_exception(max_attempts=0)(mock_func) + + with pytest.raises(RuntimeError, match="All 0 attempts failed"): + wrapped() + + mock_func.assert_not_called() + + +# --------------------------------------------------------------------------- +# async_retry_on_exception +# --------------------------------------------------------------------------- + + [email protected] +async def test_async_retry_on_exception_succeeds_first_try() -> None: + """Should await the wrapped coroutine once when it succeeds immediately.""" + mock_func = AsyncMock(return_value="ok") + wrapped = async_retry_on_exception(max_attempts=3)(mock_func) + + result = await wrapped() + + assert result == "ok" + assert mock_func.await_count == 1 + + [email protected] +async def test_async_retry_on_exception_retries_then_succeeds() -> None: + """Should retry a failing coroutine on retryable exceptions until success.""" + mock_func = AsyncMock(side_effect=[ConnectionError("fail"), "ok"]) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock) as mock_sleep: + wrapped = async_retry_on_exception( + max_attempts=3, base_delay=0.01, jitter=False + )(mock_func) + result = await wrapped() + + assert result == "ok" + assert mock_func.await_count == 2 + mock_sleep.assert_awaited_once() + + [email protected] +async def test_async_retry_on_exception_exhausts_retries_and_raises() -> None: + """Should raise the last retryable exception once max_attempts is reached.""" + mock_func = AsyncMock(side_effect=ConnectionError("nope")) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock): + wrapped = async_retry_on_exception( + max_attempts=2, base_delay=0.01, jitter=False + )(mock_func) + with pytest.raises(ConnectionError, match="nope"): + await wrapped() + + assert mock_func.await_count == 2 + + [email protected] +async def test_async_retry_on_exception_non_retryable_fails_immediately() -> None: + """Should not retry exceptions outside the configured retryable tuple.""" + mock_func = AsyncMock(side_effect=ValueError("bad")) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock) as mock_sleep: + wrapped = async_retry_on_exception( + max_attempts=3, exceptions=(ConnectionError,) + )(mock_func) + with pytest.raises(ValueError, match="bad"): + await wrapped() + + assert mock_func.await_count == 1 + mock_sleep.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# RetryableOperation +# --------------------------------------------------------------------------- + + +def test_retryable_operation_success_does_not_suppress_or_retry() -> None: + """Should leave attempt count untouched when the block succeeds.""" + op = RetryableOperation("test_op", max_attempts=3) + + with op: + pass + + assert op.current_attempt == 0 + assert op.last_exception is None + + +def test_retryable_operation_retries_until_success() -> None: + """Should suppress retryable exceptions and allow the loop to retry.""" + op = RetryableOperation("test_op", max_attempts=3, base_delay=0.01, jitter=False) + attempts = 0 Review Comment: Added the explicit `int` annotation to the attempt counter in 8ab6108ae9. ########## tests/unit_tests/mcp_service/utils/test_retry_utils.py: ########## @@ -0,0 +1,396 @@ +# 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 retry utilities. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from sqlalchemy.exc import OperationalError + +from superset.mcp_service.utils.retry_utils import ( + async_retry_database_operation, + async_retry_on_exception, + exponential_backoff, + retry_database_operation, + retry_on_exception, + retry_screenshot_operation, + RetryableOperation, +) + +SLEEP = "superset.mcp_service.utils.retry_utils.time.sleep" +ASYNC_SLEEP = "superset.mcp_service.utils.retry_utils.asyncio.sleep" + + +# --------------------------------------------------------------------------- +# exponential_backoff +# --------------------------------------------------------------------------- + + +def test_exponential_backoff_without_jitter_doubles_each_attempt() -> None: + """Should double the delay for each subsequent attempt when jitter is off.""" + assert exponential_backoff(0, base_delay=1.0, jitter=False) == 1.0 + assert exponential_backoff(1, base_delay=1.0, jitter=False) == 2.0 + assert exponential_backoff(2, base_delay=1.0, jitter=False) == 4.0 + + +def test_exponential_backoff_caps_at_max_delay() -> None: + """Should never return a delay larger than max_delay.""" + delay = exponential_backoff(10, base_delay=1.0, max_delay=5.0, jitter=False) + assert delay == 5.0 + + +def test_exponential_backoff_jitter_stays_within_twenty_five_percent() -> None: + """Should keep jittered delays within +/-25% of the base delay.""" + for _ in range(50): + delay = exponential_backoff(2, base_delay=1.0, jitter=True) + # Unjittered delay for attempt=2, base_delay=1.0 is 4.0. + assert 3.0 <= delay <= 5.0 + + +def test_exponential_backoff_never_negative() -> None: + """Should clamp the delay to be non-negative even with jitter applied.""" + for _ in range(50): + delay = exponential_backoff(0, base_delay=0.001, jitter=True) + assert delay >= 0 + + +# --------------------------------------------------------------------------- +# retry_on_exception (sync) +# --------------------------------------------------------------------------- + + +def test_retry_on_exception_succeeds_first_try() -> None: + """Should call the wrapped function once when it succeeds immediately.""" + mock_func = MagicMock(return_value="ok") + wrapped = retry_on_exception(max_attempts=3)(mock_func) + + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 1 + + +def test_retry_on_exception_retries_then_succeeds() -> None: + """Should retry on retryable exceptions and return the eventual success.""" + mock_func = MagicMock( + side_effect=[ConnectionError("fail"), ConnectionError("fail"), "ok"] + ) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, base_delay=0.01, jitter=False)( + mock_func + ) + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 3 + assert mock_sleep.call_count == 2 + + +def test_retry_on_exception_exhausts_retries_and_raises_last_exception() -> None: + """Should raise the last retryable exception once max_attempts is reached.""" + mock_func = MagicMock(side_effect=ConnectionError("always fails")) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, base_delay=0.01, jitter=False)( + mock_func + ) + with pytest.raises(ConnectionError, match="always fails"): + wrapped() + + assert mock_func.call_count == 3 + # No sleep after the final (failed) attempt. + assert mock_sleep.call_count == 2 + + +def test_retry_on_exception_non_retryable_exception_fails_immediately() -> None: + """Should not retry exceptions outside the configured retryable tuple.""" + mock_func = MagicMock(side_effect=ValueError("bad input")) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, exceptions=(ConnectionError,))( + mock_func + ) + with pytest.raises(ValueError, match="bad input"): + wrapped() + + assert mock_func.call_count == 1 + mock_sleep.assert_not_called() + + +def test_retry_on_exception_respects_custom_exceptions_tuple() -> None: + """Should retry only on the exception types passed via `exceptions`.""" + mock_func = MagicMock(side_effect=[TimeoutError("timed out"), "ok"]) + mock_func.__name__ = "mock_func" + + with patch(SLEEP): + wrapped = retry_on_exception( + max_attempts=2, base_delay=0.01, jitter=False, exceptions=(TimeoutError,) + )(mock_func) + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 2 + + +def test_retry_on_exception_max_attempts_zero_raises_runtime_error() -> None: + """Should raise RuntimeError when max_attempts leaves no attempt to run.""" + mock_func = MagicMock(return_value="ok") + mock_func.__name__ = "mock_func" + wrapped = retry_on_exception(max_attempts=0)(mock_func) + + with pytest.raises(RuntimeError, match="All 0 attempts failed"): + wrapped() + + mock_func.assert_not_called() + + +# --------------------------------------------------------------------------- +# async_retry_on_exception +# --------------------------------------------------------------------------- + + [email protected] +async def test_async_retry_on_exception_succeeds_first_try() -> None: + """Should await the wrapped coroutine once when it succeeds immediately.""" + mock_func = AsyncMock(return_value="ok") + wrapped = async_retry_on_exception(max_attempts=3)(mock_func) + + result = await wrapped() + + assert result == "ok" + assert mock_func.await_count == 1 + + [email protected] +async def test_async_retry_on_exception_retries_then_succeeds() -> None: + """Should retry a failing coroutine on retryable exceptions until success.""" + mock_func = AsyncMock(side_effect=[ConnectionError("fail"), "ok"]) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock) as mock_sleep: + wrapped = async_retry_on_exception( + max_attempts=3, base_delay=0.01, jitter=False + )(mock_func) + result = await wrapped() + + assert result == "ok" + assert mock_func.await_count == 2 + mock_sleep.assert_awaited_once() + + [email protected] +async def test_async_retry_on_exception_exhausts_retries_and_raises() -> None: + """Should raise the last retryable exception once max_attempts is reached.""" + mock_func = AsyncMock(side_effect=ConnectionError("nope")) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock): + wrapped = async_retry_on_exception( + max_attempts=2, base_delay=0.01, jitter=False + )(mock_func) + with pytest.raises(ConnectionError, match="nope"): + await wrapped() + + assert mock_func.await_count == 2 + + [email protected] +async def test_async_retry_on_exception_non_retryable_fails_immediately() -> None: + """Should not retry exceptions outside the configured retryable tuple.""" + mock_func = AsyncMock(side_effect=ValueError("bad")) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock) as mock_sleep: + wrapped = async_retry_on_exception( + max_attempts=3, exceptions=(ConnectionError,) + )(mock_func) + with pytest.raises(ValueError, match="bad"): + await wrapped() + + assert mock_func.await_count == 1 + mock_sleep.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# RetryableOperation +# --------------------------------------------------------------------------- + + +def test_retryable_operation_success_does_not_suppress_or_retry() -> None: + """Should leave attempt count untouched when the block succeeds.""" + op = RetryableOperation("test_op", max_attempts=3) + + with op: + pass + + assert op.current_attempt == 0 + assert op.last_exception is None + + +def test_retryable_operation_retries_until_success() -> None: + """Should suppress retryable exceptions and allow the loop to retry.""" + op = RetryableOperation("test_op", max_attempts=3, base_delay=0.01, jitter=False) + attempts = 0 + result = None Review Comment: Added the explicit `str | None` annotation to the nullable result in 8ab6108ae9. ########## tests/unit_tests/mcp_service/utils/test_cache_utils.py: ########## @@ -0,0 +1,268 @@ +# 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 cache utilities. +""" + +from datetime import datetime, timedelta, timezone +from typing import Any + +from superset.mcp_service.utils.cache_utils import ( + apply_cache_control_to_query_context, + get_cache_key_info, + get_cache_status_from_result, + should_use_metadata_cache, +) + +# --------------------------------------------------------------------------- +# get_cache_status_from_result +# --------------------------------------------------------------------------- + + +def test_get_cache_status_from_result_reads_from_queries_list() -> None: + """Should pull cache_hit from the first entry of a 'queries' list.""" + result: dict[str, Any] = {"queries": [{"is_cached": True}]} + Review Comment: Added the explicit `CacheStatus` annotation in 8ab6108ae9. ########## tests/unit_tests/mcp_service/utils/test_cache_utils.py: ########## @@ -0,0 +1,268 @@ +# 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 cache utilities. +""" + +from datetime import datetime, timedelta, timezone +from typing import Any + +from superset.mcp_service.utils.cache_utils import ( + apply_cache_control_to_query_context, + get_cache_key_info, + get_cache_status_from_result, + should_use_metadata_cache, +) + +# --------------------------------------------------------------------------- +# get_cache_status_from_result +# --------------------------------------------------------------------------- + + +def test_get_cache_status_from_result_reads_from_queries_list() -> None: + """Should pull cache_hit from the first entry of a 'queries' list.""" + result: dict[str, Any] = {"queries": [{"is_cached": True}]} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is True + assert status.cache_type == "query" + + +def test_get_cache_status_from_result_falls_back_to_top_level_dict() -> None: + """Should read cache info directly from the result when no 'queries' key.""" + result: dict[str, Any] = {"is_cached": False} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_empty_queries_list_uses_top_level() -> None: + """Should fall back to the top-level dict when 'queries' is an empty list.""" + result: dict[str, Any] = {"queries": [], "is_cached": True} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is True + + +def test_get_cache_status_from_result_defaults_cache_hit_to_false() -> None: + """Should treat a missing 'is_cached' key as not cached.""" + result: dict[str, Any] = {} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_reflects_force_refresh_flag() -> None: + """Should surface the caller-provided force_refresh value on the status.""" + result: dict[str, Any] = {"is_cached": True} + + status = get_cache_status_from_result(result, force_refresh=True) + + assert status.refreshed is True + + +def test_get_cache_status_from_result_parses_iso_string_cache_age() -> None: + """Should compute cache_age_seconds from an ISO-formatted cache_dttm string.""" + cache_dt = datetime.now(timezone.utc) - timedelta(seconds=120) Review Comment: Added the explicit `datetime` annotation in 8ab6108ae9. ########## tests/unit_tests/mcp_service/utils/test_cache_utils.py: ########## @@ -0,0 +1,268 @@ +# 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 cache utilities. +""" + +from datetime import datetime, timedelta, timezone +from typing import Any + +from superset.mcp_service.utils.cache_utils import ( + apply_cache_control_to_query_context, + get_cache_key_info, + get_cache_status_from_result, + should_use_metadata_cache, +) + +# --------------------------------------------------------------------------- +# get_cache_status_from_result +# --------------------------------------------------------------------------- + + +def test_get_cache_status_from_result_reads_from_queries_list() -> None: + """Should pull cache_hit from the first entry of a 'queries' list.""" + result: dict[str, Any] = {"queries": [{"is_cached": True}]} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is True + assert status.cache_type == "query" + + +def test_get_cache_status_from_result_falls_back_to_top_level_dict() -> None: + """Should read cache info directly from the result when no 'queries' key.""" + result: dict[str, Any] = {"is_cached": False} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_empty_queries_list_uses_top_level() -> None: + """Should fall back to the top-level dict when 'queries' is an empty list.""" + result: dict[str, Any] = {"queries": [], "is_cached": True} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is True + + +def test_get_cache_status_from_result_defaults_cache_hit_to_false() -> None: + """Should treat a missing 'is_cached' key as not cached.""" + result: dict[str, Any] = {} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_reflects_force_refresh_flag() -> None: + """Should surface the caller-provided force_refresh value on the status.""" + result: dict[str, Any] = {"is_cached": True} + + status = get_cache_status_from_result(result, force_refresh=True) + + assert status.refreshed is True + + +def test_get_cache_status_from_result_parses_iso_string_cache_age() -> None: + """Should compute cache_age_seconds from an ISO-formatted cache_dttm string.""" + cache_dt = datetime.now(timezone.utc) - timedelta(seconds=120) + result: dict[str, Any] = { + "is_cached": True, + "cache_dttm": cache_dt.isoformat(), + } + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds >= 119 + + +def test_get_cache_status_from_result_parses_z_suffixed_string() -> None: + """Should handle a trailing 'Z' UTC designator in the cache_dttm string.""" + result: dict[str, Any] = { + "is_cached": True, + "cache_dttm": "2020-01-01T00:00:00Z", + } + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds > 0 + + +def test_get_cache_status_from_result_parses_datetime_object_cache_age() -> None: + """Should compute cache_age_seconds when cache_dttm is already a datetime.""" + cache_dt = datetime.now(timezone.utc) - timedelta(seconds=60) + result: dict[str, Any] = {"is_cached": True, "cache_dttm": cache_dt} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds >= 59 + + +def test_get_cache_status_from_result_handles_unparseable_cache_dttm() -> None: + """Should swallow parse errors and leave cache_age_seconds as None.""" + result: dict[str, Any] = {"is_cached": True, "cache_dttm": "not-a-date"} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is None + # The rest of the status should still be populated correctly. + assert status.cache_hit is True + + +def test_get_cache_status_from_result_no_cache_dttm_leaves_age_none() -> None: + """Should leave cache_age_seconds as None when cache_dttm is absent.""" + result: dict[str, Any] = {"is_cached": True} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is None + + +# --------------------------------------------------------------------------- +# apply_cache_control_to_query_context +# --------------------------------------------------------------------------- + + +def test_apply_cache_control_sets_force_when_cache_disabled() -> None: + """Should set force=True on the query context when use_cache is False.""" + ctx = apply_cache_control_to_query_context({"queries": []}, use_cache=False) Review Comment: Added the explicit `dict[str, Any]` annotation in 8ab6108ae9. ########## tests/unit_tests/mcp_service/utils/test_cache_utils.py: ########## @@ -0,0 +1,268 @@ +# 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 cache utilities. +""" + +from datetime import datetime, timedelta, timezone +from typing import Any + +from superset.mcp_service.utils.cache_utils import ( + apply_cache_control_to_query_context, + get_cache_key_info, + get_cache_status_from_result, + should_use_metadata_cache, +) + +# --------------------------------------------------------------------------- +# get_cache_status_from_result +# --------------------------------------------------------------------------- + + +def test_get_cache_status_from_result_reads_from_queries_list() -> None: + """Should pull cache_hit from the first entry of a 'queries' list.""" + result: dict[str, Any] = {"queries": [{"is_cached": True}]} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is True + assert status.cache_type == "query" + + +def test_get_cache_status_from_result_falls_back_to_top_level_dict() -> None: + """Should read cache info directly from the result when no 'queries' key.""" + result: dict[str, Any] = {"is_cached": False} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_empty_queries_list_uses_top_level() -> None: + """Should fall back to the top-level dict when 'queries' is an empty list.""" + result: dict[str, Any] = {"queries": [], "is_cached": True} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is True + + +def test_get_cache_status_from_result_defaults_cache_hit_to_false() -> None: + """Should treat a missing 'is_cached' key as not cached.""" + result: dict[str, Any] = {} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_reflects_force_refresh_flag() -> None: + """Should surface the caller-provided force_refresh value on the status.""" + result: dict[str, Any] = {"is_cached": True} + + status = get_cache_status_from_result(result, force_refresh=True) + + assert status.refreshed is True + + +def test_get_cache_status_from_result_parses_iso_string_cache_age() -> None: + """Should compute cache_age_seconds from an ISO-formatted cache_dttm string.""" + cache_dt = datetime.now(timezone.utc) - timedelta(seconds=120) + result: dict[str, Any] = { + "is_cached": True, + "cache_dttm": cache_dt.isoformat(), + } + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds >= 119 + + +def test_get_cache_status_from_result_parses_z_suffixed_string() -> None: + """Should handle a trailing 'Z' UTC designator in the cache_dttm string.""" + result: dict[str, Any] = { + "is_cached": True, + "cache_dttm": "2020-01-01T00:00:00Z", + } + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds > 0 + + +def test_get_cache_status_from_result_parses_datetime_object_cache_age() -> None: + """Should compute cache_age_seconds when cache_dttm is already a datetime.""" + cache_dt = datetime.now(timezone.utc) - timedelta(seconds=60) + result: dict[str, Any] = {"is_cached": True, "cache_dttm": cache_dt} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds >= 59 + + +def test_get_cache_status_from_result_handles_unparseable_cache_dttm() -> None: + """Should swallow parse errors and leave cache_age_seconds as None.""" + result: dict[str, Any] = {"is_cached": True, "cache_dttm": "not-a-date"} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is None + # The rest of the status should still be populated correctly. + assert status.cache_hit is True + + +def test_get_cache_status_from_result_no_cache_dttm_leaves_age_none() -> None: + """Should leave cache_age_seconds as None when cache_dttm is absent.""" + result: dict[str, Any] = {"is_cached": True} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is None + + +# --------------------------------------------------------------------------- +# apply_cache_control_to_query_context +# --------------------------------------------------------------------------- + + +def test_apply_cache_control_sets_force_when_cache_disabled() -> None: + """Should set force=True on the query context when use_cache is False.""" + ctx = apply_cache_control_to_query_context({"queries": []}, use_cache=False) + + assert ctx["force"] is True + + +def test_apply_cache_control_sets_force_when_force_refresh_requested() -> None: + """Should set force=True on the query context when force_refresh is True.""" + ctx = apply_cache_control_to_query_context( + {"queries": []}, use_cache=True, force_refresh=True + ) + + assert ctx["force"] is True + + +def test_apply_cache_control_does_not_set_force_when_using_cache() -> None: + """Should leave 'force' unset when caching is enabled and not refreshing.""" + ctx = apply_cache_control_to_query_context( + {"queries": []}, use_cache=True, force_refresh=False + ) + + assert "force" not in ctx + + +def test_apply_cache_control_applies_cache_timeout_to_every_query() -> None: + """Should stamp cache_timeout onto every query in the context.""" + ctx = apply_cache_control_to_query_context( + {"queries": [{"metric": "a"}, {"metric": "b"}]}, cache_timeout=60 + ) + + assert ctx["queries"][0]["cache_timeout"] == 60 + assert ctx["queries"][1]["cache_timeout"] == 60 + + +def test_apply_cache_control_leaves_queries_untouched_when_timeout_is_none() -> None: + """Should not add a cache_timeout key when none is provided.""" + ctx = apply_cache_control_to_query_context( + {"queries": [{"metric": "a"}]}, cache_timeout=None + ) + + assert "cache_timeout" not in ctx["queries"][0] + + +def test_apply_cache_control_missing_queries_key_is_a_noop_for_timeout() -> None: + """Should not raise or add a 'queries' key when applying a cache_timeout.""" + ctx = apply_cache_control_to_query_context({}, cache_timeout=60) + + assert "queries" not in ctx + + +def test_apply_cache_control_returns_the_same_dict_instance() -> None: + """Should mutate and return the same query_context object it was given.""" + original: dict[str, Any] = {"queries": []} + + result = apply_cache_control_to_query_context(original, use_cache=False) + + assert result is original + + +# --------------------------------------------------------------------------- +# should_use_metadata_cache +# --------------------------------------------------------------------------- + + +def test_should_use_metadata_cache_true_by_default() -> None: + """Should default to using the metadata cache.""" + assert should_use_metadata_cache() is True + + +def test_should_use_metadata_cache_false_when_cache_disabled() -> None: + """Should return False when use_cache is False.""" + assert should_use_metadata_cache(use_cache=False) is False + + +def test_should_use_metadata_cache_false_when_refresh_requested() -> None: + """Should return False when refresh_metadata is True, even with use_cache.""" + assert should_use_metadata_cache(use_cache=True, refresh_metadata=True) is False + + +def test_should_use_metadata_cache_true_when_enabled_and_not_refreshing() -> None: + """Should return True when caching is enabled and no refresh is requested.""" + assert should_use_metadata_cache(use_cache=True, refresh_metadata=False) is True + + +# --------------------------------------------------------------------------- +# get_cache_key_info +# --------------------------------------------------------------------------- + + +def test_get_cache_key_info_returns_none_for_none_input() -> None: + """Should return None when no cache key is provided.""" + assert get_cache_key_info(None) is None + + +def test_get_cache_key_info_returns_none_for_empty_string() -> None: + """Should treat an empty string cache key the same as no key.""" + assert get_cache_key_info("") is None + + +def test_get_cache_key_info_returns_short_keys_unchanged() -> None: + """Should return short cache keys unmodified.""" + key = "short_cache_key" Review Comment: Added the explicit `str` annotation in 8ab6108ae9. ########## tests/unit_tests/mcp_service/database/tool/test_database_tools.py: ########## @@ -120,7 +121,7 @@ def mock_auth(): @pytest.fixture(autouse=True) -def allow_data_model_metadata(): +def _allow_data_model_metadata(): Review Comment: Added the explicit `Iterator[None]` fixture return annotation in 8ab6108ae9. ########## tests/unit_tests/mcp_service/semantic_layer/tool/test_get_table.py: ########## @@ -297,3 +298,227 @@ async def test_get_table_external_time_range_without_dttm_validation_error( assert data["success"] is False assert data["error_type"] == "ValidationError" assert "no datetime dimension" in data["message"] + + [email protected]() +async def test_get_table_dataset_not_found(mcp_server: FastMCP) -> None: + """get_table returns NotFound when dataset_id doesn't resolve to a dataset.""" + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=None): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + {"request": {"dataset_id": 999999, "metrics": ["revenue"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "NotFound" + assert "999999" in data["message"] + + [email protected]() +async def test_get_table_view_not_found(mcp_server: FastMCP) -> None: + """get_table returns NotFound when view_id doesn't resolve to a view.""" + with patch( + "superset.daos.semantic_layer.SemanticViewDAO.find_by_id", + return_value=None, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + {"request": {"view_id": 999999, "metrics": ["bookings"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "NotFound" + assert "999999" in data["message"] + + [email protected]() +async def test_get_table_invalid_filter_column_validation_error( + mcp_server: FastMCP, +) -> None: + """get_table errors when a filter references an unknown column.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "filters": [{"col": "bogus_col", "op": "==", "val": "x"}], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown filter column: 'bogus_col'" in data["error"] + + [email protected]() +async def test_get_table_invalid_order_by_validation_error( + mcp_server: FastMCP, +) -> None: + """get_table errors when order_by references an unknown column/metric.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "order_by": ["bogus_order_col"], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown order_by: 'bogus_order_col'" in data["error"] + + [email protected]() +async def test_get_table_unknown_filter_operator_passes_through( + mcp_server: FastMCP, +) -> None: + """An operator string outside the documented set is not schema-validated. + + ``GetTableFilter.op`` is a plain ``str`` field (not a Literal/Enum), so + the tool does not reject unrecognized operator values itself -- it + forwards them verbatim to the query layer, which is responsible for + interpreting/rejecting them. + """ + mock_ds = _make_dataset(42) + query_result: dict[str, Any] = { + "queries": [ + { + "data": [{"region": "west", "revenue": 100}], + "colnames": ["region", "revenue"], + "rowcount": 1, + } + ] + } + + with ( + patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds), + patch( + "superset.commands.chart.data.get_data_command.ChartDataCommand" + ) as mock_command_cls, + patch( + "superset.common.query_context_factory.QueryContextFactory" + ) as mock_factory_cls, + ): + mock_command_cls.return_value.run.return_value = query_result + mock_factory_cls.return_value.create.return_value = MagicMock() + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "filters": [ + {"col": "region", "op": "TOTALLY_BOGUS_OP", "val": "x"} + ], + } + }, + ) + data = json.loads(result.content[0].text) + + create_kwargs = mock_factory_cls.return_value.create.call_args.kwargs Review Comment: Added the explicit `dict[str, Any]` kwargs annotation in 8ab6108ae9. ########## tests/unit_tests/mcp_service/semantic_layer/tool/test_get_table.py: ########## @@ -297,3 +298,227 @@ async def test_get_table_external_time_range_without_dttm_validation_error( assert data["success"] is False assert data["error_type"] == "ValidationError" assert "no datetime dimension" in data["message"] + + [email protected]() +async def test_get_table_dataset_not_found(mcp_server: FastMCP) -> None: + """get_table returns NotFound when dataset_id doesn't resolve to a dataset.""" + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=None): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + {"request": {"dataset_id": 999999, "metrics": ["revenue"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "NotFound" + assert "999999" in data["message"] + + [email protected]() +async def test_get_table_view_not_found(mcp_server: FastMCP) -> None: + """get_table returns NotFound when view_id doesn't resolve to a view.""" + with patch( + "superset.daos.semantic_layer.SemanticViewDAO.find_by_id", + return_value=None, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + {"request": {"view_id": 999999, "metrics": ["bookings"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "NotFound" + assert "999999" in data["message"] + + [email protected]() +async def test_get_table_invalid_filter_column_validation_error( + mcp_server: FastMCP, +) -> None: + """get_table errors when a filter references an unknown column.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "filters": [{"col": "bogus_col", "op": "==", "val": "x"}], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown filter column: 'bogus_col'" in data["error"] + + [email protected]() +async def test_get_table_invalid_order_by_validation_error( + mcp_server: FastMCP, +) -> None: + """get_table errors when order_by references an unknown column/metric.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "order_by": ["bogus_order_col"], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown order_by: 'bogus_order_col'" in data["error"] + + [email protected]() +async def test_get_table_unknown_filter_operator_passes_through( + mcp_server: FastMCP, +) -> None: + """An operator string outside the documented set is not schema-validated. + + ``GetTableFilter.op`` is a plain ``str`` field (not a Literal/Enum), so + the tool does not reject unrecognized operator values itself -- it + forwards them verbatim to the query layer, which is responsible for + interpreting/rejecting them. + """ + mock_ds = _make_dataset(42) + query_result: dict[str, Any] = { + "queries": [ + { + "data": [{"region": "west", "revenue": 100}], + "colnames": ["region", "revenue"], + "rowcount": 1, + } + ] + } + + with ( + patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds), + patch( + "superset.commands.chart.data.get_data_command.ChartDataCommand" + ) as mock_command_cls, + patch( + "superset.common.query_context_factory.QueryContextFactory" + ) as mock_factory_cls, + ): + mock_command_cls.return_value.run.return_value = query_result + mock_factory_cls.return_value.create.return_value = MagicMock() + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "filters": [ + {"col": "region", "op": "TOTALLY_BOGUS_OP", "val": "x"} + ], + } + }, + ) + data = json.loads(result.content[0].text) + + create_kwargs = mock_factory_cls.return_value.create.call_args.kwargs + forwarded_filters = create_kwargs["queries"][0]["filters"] Review Comment: Added the explicit `list[dict[str, Any]]` filters annotation in 8ab6108ae9. -- 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]
