|
3 | 3 | from unittest.mock import MagicMock, patch |
4 | 4 |
|
5 | 5 | from spark_history_mcp.api.spark_client import SparkRestClient |
| 6 | +from spark_history_mcp.config.config import ServerConfig |
6 | 7 | from spark_history_mcp.models.spark_types import ( |
7 | 8 | ApplicationInfo, |
8 | 9 | ExecutionData, |
@@ -1087,3 +1088,73 @@ def test_get_slowest_sql_queries_limit(self, mock_get_client): |
1087 | 1088 | self.assertEqual(result[0].duration, 10000) |
1088 | 1089 | self.assertEqual(result[1].duration, 9000) |
1089 | 1090 | self.assertEqual(result[2].duration, 8000) |
| 1091 | + |
| 1092 | + @patch("spark_history_mcp.tools.tools.get_client_or_default") |
| 1093 | + def test_list_slowest_sql_queries_uses_server_config_for_plan_description( |
| 1094 | + self, mock_get_client |
| 1095 | + ): |
| 1096 | + """Test that include_plan_description falls back to server config when not provided""" |
| 1097 | + # Setup mock client with server config |
| 1098 | + mock_client = MagicMock() |
| 1099 | + server_config = ServerConfig( |
| 1100 | + url="http://test:18080", include_plan_description=False |
| 1101 | + ) |
| 1102 | + mock_client.config = server_config |
| 1103 | + |
| 1104 | + # Create mock SQL execution |
| 1105 | + sql = MagicMock(spec=ExecutionData) |
| 1106 | + sql.id = 1 |
| 1107 | + sql.duration = 5000 |
| 1108 | + sql.status = "COMPLETED" |
| 1109 | + sql.success_job_ids = [1] |
| 1110 | + sql.failed_job_ids = [] |
| 1111 | + sql.running_job_ids = [] |
| 1112 | + sql.description = "Test Query" |
| 1113 | + sql.submission_time = datetime.now() |
| 1114 | + sql.plan_description = "Sample plan description" |
| 1115 | + |
| 1116 | + mock_client.get_sql_list.return_value = [sql] |
| 1117 | + mock_get_client.return_value = mock_client |
| 1118 | + |
| 1119 | + # Call function without include_plan_description parameter (should use server config) |
| 1120 | + result = list_slowest_sql_queries("spark-app-123") |
| 1121 | + |
| 1122 | + # Verify plan description is empty due to server config setting False |
| 1123 | + self.assertEqual(len(result), 1) |
| 1124 | + self.assertEqual(result[0].plan_description, "") |
| 1125 | + |
| 1126 | + @patch("spark_history_mcp.tools.tools.get_client_or_default") |
| 1127 | + def test_list_slowest_sql_queries_explicit_override_server_config( |
| 1128 | + self, mock_get_client |
| 1129 | + ): |
| 1130 | + """Test that server config overrides parameter when config is set""" |
| 1131 | + # Setup mock client with server config set to False |
| 1132 | + mock_client = MagicMock() |
| 1133 | + server_config = ServerConfig( |
| 1134 | + url="http://test:18080", include_plan_description=False |
| 1135 | + ) |
| 1136 | + mock_client.config = server_config |
| 1137 | + |
| 1138 | + # Create mock SQL execution |
| 1139 | + sql = MagicMock(spec=ExecutionData) |
| 1140 | + sql.id = 1 |
| 1141 | + sql.duration = 5000 |
| 1142 | + sql.status = "COMPLETED" |
| 1143 | + sql.success_job_ids = [1] |
| 1144 | + sql.failed_job_ids = [] |
| 1145 | + sql.running_job_ids = [] |
| 1146 | + sql.description = "Test Query" |
| 1147 | + sql.submission_time = datetime.now() |
| 1148 | + sql.plan_description = "Sample plan description" |
| 1149 | + |
| 1150 | + mock_client.get_sql_list.return_value = [sql] |
| 1151 | + mock_get_client.return_value = mock_client |
| 1152 | + |
| 1153 | + # Call function with explicit include_plan_description=True (config should override to False) |
| 1154 | + result = list_slowest_sql_queries( |
| 1155 | + "spark-app-123", include_plan_description=True |
| 1156 | + ) |
| 1157 | + |
| 1158 | + # Verify plan description is NOT included because server config is False |
| 1159 | + self.assertEqual(len(result), 1) |
| 1160 | + self.assertEqual(result[0].plan_description, "") |
0 commit comments