From 0303a234a3ed61390b0fb1f8deea111e94aa3219 Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Tue, 10 Mar 2026 12:30:28 -0400 Subject: [PATCH] Fix tests --- superset/semantic_layers/mapper.py | 22 +++------------- .../unit_tests/semantic_layers/mapper_test.py | 26 +++++++++---------- 2 files changed, 16 insertions(+), 32 deletions(-) diff --git a/superset/semantic_layers/mapper.py b/superset/semantic_layers/mapper.py index d3220154cc0..66c86cf6d4e 100644 --- a/superset/semantic_layers/mapper.py +++ b/superset/semantic_layers/mapper.py @@ -115,7 +115,7 @@ def get_results(query_object: QueryObject) -> QueryResult: dispatcher = ( semantic_view.get_row_count if query_object.is_rowcount - else semantic_view.get_dataframe + else semantic_view.get_table ) # Step 1: Convert QueryObject to list of SemanticQuery objects @@ -124,15 +124,7 @@ def get_results(query_object: QueryObject) -> QueryResult: # Step 2: Execute the main query (first in the list) main_query = queries[0] - main_result = dispatcher( - metrics=main_query.metrics, - dimensions=main_query.dimensions, - filters=main_query.filters, - order=main_query.order, - limit=main_query.limit, - offset=main_query.offset, - group_limit=main_query.group_limit, - ) + main_result = dispatcher(main_query) main_df = main_result.results.to_pandas() @@ -163,15 +155,7 @@ def get_results(query_object: QueryObject) -> QueryResult: strict=False, ): # Execute the offset query - result = dispatcher( - metrics=offset_query.metrics, - dimensions=offset_query.dimensions, - filters=offset_query.filters, - order=offset_query.order, - limit=offset_query.limit, - offset=offset_query.offset, - group_limit=offset_query.group_limit, - ) + result = dispatcher(offset_query) # Add this query's requests to the collection all_requests.extend(result.requests) diff --git a/tests/unit_tests/semantic_layers/mapper_test.py b/tests/unit_tests/semantic_layers/mapper_test.py index 37711ee1686..1d9028ee41f 100644 --- a/tests/unit_tests/semantic_layers/mapper_test.py +++ b/tests/unit_tests/semantic_layers/mapper_test.py @@ -1228,7 +1228,7 @@ def test_get_results_without_time_offsets( } ) - # Mock the semantic view's get_dataframe method + # Mock the semantic view's get_table method mock_result = SemanticResult( requests=[ SemanticRequest( @@ -1239,7 +1239,7 @@ def test_get_results_without_time_offsets( results=pa.Table.from_pandas(main_df), ) - mock_datasource.implementation.get_dataframe = mocker.Mock(return_value=mock_result) + mock_datasource.implementation.get_table = mocker.Mock(return_value=mock_result) # Create query object without time offsets query_object = ValidatedQueryObject( @@ -1284,7 +1284,7 @@ def test_get_results_with_single_time_offset( } ) - # Mock the semantic view's get_dataframe method + # Mock the semantic view's get_table method # It will be called twice: once for main, once for offset mock_main_result = SemanticResult( requests=[ @@ -1312,7 +1312,7 @@ def test_get_results_with_single_time_offset( results=pa.Table.from_pandas(offset_df.copy()), ) - mock_datasource.implementation.get_dataframe = mocker.Mock( + mock_datasource.implementation.get_table = mocker.Mock( side_effect=[mock_main_result, mock_offset_result] ) @@ -1391,7 +1391,7 @@ def test_get_results_with_multiple_time_offsets( results=pa.Table.from_pandas(offset_1m_df.copy()), ) - mock_datasource.implementation.get_dataframe = mocker.Mock( + mock_datasource.implementation.get_table = mocker.Mock( side_effect=[mock_main_result, mock_offset_1w_result, mock_offset_1m_result] ) @@ -1457,7 +1457,7 @@ def test_get_results_with_empty_offset_result( results=pa.Table.from_pandas(offset_df), ) - mock_datasource.implementation.get_dataframe = mocker.Mock( + mock_datasource.implementation.get_table = mocker.Mock( side_effect=[mock_main_result, mock_offset_result] ) @@ -1519,7 +1519,7 @@ def test_get_results_with_partial_offset_match( results=pa.Table.from_pandas(offset_df.copy()), ) - mock_datasource.implementation.get_dataframe = mocker.Mock( + mock_datasource.implementation.get_table = mocker.Mock( side_effect=[mock_main_result, mock_offset_result] ) @@ -1584,7 +1584,7 @@ def test_get_results_with_multiple_dimensions( results=pa.Table.from_pandas(offset_df.copy()), ) - mock_datasource.implementation.get_dataframe = mocker.Mock( + mock_datasource.implementation.get_table = mocker.Mock( side_effect=[mock_main_result, mock_offset_result] ) @@ -1663,7 +1663,7 @@ def test_get_results_with_duplicate_columns( results=pa.Table.from_pandas(offset_df.copy()), ) - mock_datasource.implementation.get_dataframe = mocker.Mock( + mock_datasource.implementation.get_table = mocker.Mock( side_effect=[mock_main_result, mock_offset_result] ) @@ -1702,7 +1702,7 @@ def test_get_results_empty_requests( results=pa.Table.from_pandas(main_df), ) - mock_datasource.implementation.get_dataframe = mocker.Mock(return_value=mock_result) + mock_datasource.implementation.get_table = mocker.Mock(return_value=mock_result) query_object = ValidatedQueryObject( datasource=mock_datasource, @@ -2326,7 +2326,7 @@ def test_get_results_with_is_rowcount( ) mock_datasource.implementation.get_row_count = mocker.Mock(return_value=mock_result) - mock_datasource.implementation.get_dataframe = mocker.Mock() + mock_datasource.implementation.get_table = mocker.Mock() query_object = ValidatedQueryObject( datasource=mock_datasource, @@ -2340,9 +2340,9 @@ def test_get_results_with_is_rowcount( result = get_results(query_object) - # Should have called get_row_count, not get_dataframe + # Should have called get_row_count, not get_table mock_datasource.implementation.get_row_count.assert_called_once() - mock_datasource.implementation.get_dataframe.assert_not_called() + mock_datasource.implementation.get_table.assert_not_called() pd.testing.assert_frame_equal(result.df, main_df)