sadpandajoe commented on code in PR #44406:
URL: https://github.com/apache/superset/pull/44406#discussion_r4329621805
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -3167,11 +3915,211 @@ def payload(context: QueryContext, **kwargs: Any) ->
dict[str, Any]:
datasource="native-1", annotation_context=context
)
assert current_key != legacy_key
- assert context == {
- **legacy_context,
- "semantic_filter_protocol": "semantic-null-filters-v1",
- "source_versions": {"7": view.metadata_generation},
- }
+ assert context is not None
+ assert context["source_scope"][0]["semantic_filter_protocol"] == (
+ "semantic-null-filters-v1"
+ )
+ assert context["source_scope"][0]["metadata_version"] ==
view.metadata_generation
+
+
+def _annotated_payload_fixture(
+ *, status: str = "success"
+) -> tuple[QueryContextProcessor, Any, MagicMock]:
+ """A processor, annotated query object, and a dataframe cache hit."""
+ from superset.common.query_object import QueryObject
+
+ mock_query_context = MagicMock()
+ mock_query_context.force = False
+ mock_datasource = MagicMock()
+ mock_datasource.column_names = ["col1"]
+ processor = QueryContextProcessor(mock_query_context)
+ processor._qc_datasource = mock_datasource
+ query_obj = QueryObject(
+ datasource=mock_datasource,
+ columns=["col1"],
+ annotation_layers=[
+ {
+ "annotationType": "EVENT",
+ "sourceType": "NATIVE",
+ "name": "a",
+ "value": 1,
+ }
+ ],
+ )
+ cache = MagicMock()
+ cache.is_loaded = True
+ cache.applied_filter_columns = ["col1"]
+ cache.df = pd.DataFrame({"col1": [1, 2, 3]})
+ cache.query = ""
+ cache.status = status
+ cache.error_message = "dataframe failed" if status == QueryStatus.FAILED
else None
+ cache.is_cached = True
+ cache.applied_template_filters = []
+ cache.rejected_filter_columns = []
+ cache.bq_memory_limited = False
+ cache.result_persisted = False
+ return processor, query_obj, cache
+
+
+def _run_annotated_payload(
+ processor: QueryContextProcessor,
+ query_obj: Any,
+ cache: MagicMock,
+ *,
+ force_cached: bool = False,
+) -> Any:
+ with (
+ patch(
+ "superset.common.query_context_processor.QueryCacheManager"
+ ) as mock_cache_manager,
+ patch.object(query_obj, "validate", return_value=None),
+ patch.object(processor, "query_cache_key", return_value="df-key"),
+ patch.object(processor, "annotation_cache_key",
return_value="ann-key"),
+ patch.object(processor, "get_cache_timeout", return_value=3600),
+ ):
+ mock_cache_manager.get.return_value = cache
+ return processor.get_df_payload_result(query_obj,
force_cached=force_cached)
+
+
+def test_get_df_payload_result_skips_annotations_when_dataframe_failed() ->
None:
+ """A failed dataframe keeps its own error and never triggers the fetch."""
+ processor, query_obj, cache =
_annotated_payload_fixture(status=QueryStatus.FAILED)
+ with patch.object(processor, "_get_annotation_data_cached") as
mock_annotation:
+ result = _run_annotated_payload(processor, query_obj, cache)
+ mock_annotation.assert_not_called()
+ assert result.payload["annotation_data"] == {}
+ assert result.payload["error"] == "dataframe failed"
+
+
+def test_get_df_payload_result_force_cached_misses_without_annotation_entry()
-> None:
+ """Through the real helper, ``force_cached`` with no annotation entry
+ raises ``CacheLoadError`` instead of running the annotation SQL."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(processor, "get_annotation_data") as mock_live,
+ ):
+ cm.data_cache.get.return_value = None
+ with pytest.raises(CacheLoadError):
+ _run_annotated_payload(processor, query_obj, cache,
force_cached=True)
+ mock_live.assert_not_called()
+
+
+def test_get_df_payload_result_annotation_failure_marks_payload_failed() ->
None:
+ """An annotation error on a good dataframe fails the payload."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(
+ processor,
+ "get_annotation_data",
+ side_effect=QueryObjectValidationError("annotation boom"),
+ ),
+ ):
+ cm.data_cache.get.return_value = None
+ result = _run_annotated_payload(processor, query_obj, cache)
+ assert cache.status == QueryStatus.FAILED
+ assert result.payload["error"] == "annotation boom"
+
+
+def test_get_df_payload_result_completeness_error_escapes_annotation_fetch()
-> None:
+ """A semantic annotation's incomplete-result signal must propagate (as the
+ dataframe path does) so the async task is not reported successful."""
+ from superset.exceptions import SemanticResultCompletenessError
+
+ processor, query_obj, cache = _annotated_payload_fixture()
+ with (
+ patch("superset.common.query_context_processor.cache_manager") as cm,
+ patch.object(
+ processor,
+ "get_annotation_data",
+ side_effect=SemanticResultCompletenessError("unverified"),
+ ),
+ ):
+ cm.data_cache.get.return_value = None
+ with pytest.raises(SemanticResultCompletenessError):
+ _run_annotated_payload(processor, query_obj, cache)
+
+
+def test_get_df_payload_result_times_annotation_fetch_as_data_acquisition() ->
None:
+ """Annotation work lands in ``data_acquisition_ns`` even on a dataframe
+ cache hit (which otherwise reports no acquisition time)."""
+ processor, query_obj, cache = _annotated_payload_fixture()
+
+ def slow_annotations(*_: Any, **__: Any) -> dict[str, Any]:
+ time.sleep(0.01)
+ return {}
+
+ with patch.object(
+ processor, "_get_annotation_data_cached", side_effect=slow_annotations
+ ):
+ result = _run_annotated_payload(processor, query_obj, cache)
+ assert result.timing.data_acquisition_ns is not None
+ assert result.timing.data_acquisition_ns >= 10_000_000
Review Comment:
The new timing test only runs with a dataframe cache hit, so the summing
path is never exercised: if the annotation duration replaced the dataframe time
instead of being added to it, a cold annotated chart would report only the
annotation time as `data_acquisition_ms` and every assertion here would still
pass. Could you add a case where both the dataframe fetch and the annotation
fetch run with a deterministic clock (say 30 ns and 20 ns) and assert
`data_acquisition_ns == 50`?
--
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]
