From 70114ce0bba7c4c47de33bdfb864b7a9814136ef Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Thu, 3 Dec 2020 16:56:22 +0200 Subject: [PATCH] add tests + refactor fixtures --- superset/charts/api.py | 1 + superset/charts/schemas.py | 5 +- superset/common/query_context.py | 2 +- superset/utils/core.py | 7 ++ tests/annotation_layers/api_tests.py | 79 +++------------------ tests/annotation_layers/fixtures.py | 101 +++++++++++++++++++++++++++ tests/charts/api_tests.py | 52 ++++++++++++-- tests/fixtures/query_context.py | 73 +++++++++++++++++++ 8 files changed, 242 insertions(+), 78 deletions(-) create mode 100644 tests/annotation_layers/fixtures.py diff --git a/superset/charts/api.py b/superset/charts/api.py index 75473e71a9072..450e5e9f6eb21 100644 --- a/superset/charts/api.py +++ b/superset/charts/api.py @@ -485,6 +485,7 @@ def data(self) -> Response: """ if request.is_json: json_body = request.json + print(json_body) elif request.form.get("form_data"): # CSV export submits regular form data json_body = json.loads(request.form["form_data"]) diff --git a/superset/charts/schemas.py b/superset/charts/schemas.py index ca1497a80e6ee..5e346fad37597 100644 --- a/superset/charts/schemas.py +++ b/superset/charts/schemas.py @@ -23,6 +23,7 @@ from superset.common.query_context import QueryContext from superset.utils import schema as utils from superset.utils.core import ( + AnnotationType, FilterOperator, PostProcessingBoxplotWhiskerType, PostProcessingContributionOrientation, @@ -783,9 +784,7 @@ class ChartDataExtrasSchema(Schema): class AnnotationLayerSchema(Schema): annotationType = fields.String( description="Type of annotation layer", - validate=validate.OneOf( - choices=("EVENT", "FORMULA", "INTERVAL", "TIME_SERIES",) - ), + validate=validate.OneOf(choices=[ann.value for ann in AnnotationType]), ) color = fields.String(description="Layer color", allow_none=True,) descriptionColumns = fields.List( diff --git a/superset/common/query_context.py b/superset/common/query_context.py index c176783835c5f..c30d7170685cf 100644 --- a/superset/common/query_context.py +++ b/superset/common/query_context.py @@ -295,7 +295,7 @@ def get_annotation_data(self, query_obj: QueryObject) -> Dict[str, Any]: ) return annotation_data - def get_df_payload( # pylint: disable=too-many-statements + def get_df_payload( # pylint: disable=too-many-statements,too-many-locals self, query_obj: QueryObject, **kwargs: Any ) -> Dict[str, Any]: """Handles caching around the df payload retrieval""" diff --git a/superset/utils/core.py b/superset/utils/core.py index b4e3ef51e12e2..2d4febeb426cb 100644 --- a/superset/utils/core.py +++ b/superset/utils/core.py @@ -1591,6 +1591,13 @@ class ExtraFiltersTimeColumnType(str, Enum): TIME_RANGE = "__time_range" +class AnnotationType(str, Enum): + FORMULA = "FORMULA" + INTERVAL = "INTERVAL" + EVENT = "EVENT" + TIME_SERIES = "TIME_SERIES" + + def is_test() -> bool: return strtobool(os.environ.get("SUPERSET_TESTENV", "false")) diff --git a/tests/annotation_layers/api_tests.py b/tests/annotation_layers/api_tests.py index 2b0df4126be26..0ee361bcfbd1b 100644 --- a/tests/annotation_layers/api_tests.py +++ b/tests/annotation_layers/api_tests.py @@ -16,8 +16,6 @@ # under the License. # isort:skip_file """Unit tests for Superset""" -from datetime import datetime -from typing import Optional import json import pytest @@ -29,77 +27,17 @@ from superset.models.annotations import Annotation, AnnotationLayer from tests.base_tests import SupersetTestCase - +from tests.annotation_layers.fixtures import ( + create_annotation_layers, + get_end_dttm, + get_start_dttm, +) ANNOTATION_LAYERS_COUNT = 10 ANNOTATIONS_COUNT = 5 class TestAnnotationLayerApi(SupersetTestCase): - def insert_annotation_layer( - self, name: str = "", descr: str = "" - ) -> AnnotationLayer: - annotation_layer = AnnotationLayer(name=name, descr=descr,) - db.session.add(annotation_layer) - db.session.commit() - return annotation_layer - - def insert_annotation( - self, - layer: AnnotationLayer, - short_descr: str, - long_descr: str, - json_metadata: Optional[str] = "", - start_dttm: Optional[datetime] = None, - end_dttm: Optional[datetime] = None, - ) -> Annotation: - annotation = Annotation( - layer=layer, - short_descr=short_descr, - long_descr=long_descr, - json_metadata=json_metadata, - start_dttm=start_dttm, - end_dttm=end_dttm, - ) - db.session.add(annotation) - db.session.commit() - return annotation - - @pytest.fixture() - def create_annotation_layers(self): - """ - Creates ANNOTATION_LAYERS_COUNT-1 layers with no annotations - and a final one with ANNOTATION_COUNT childs - :return: - """ - with self.create_app().app_context(): - annotation_layers = [] - annotations = [] - for cx in range(ANNOTATION_LAYERS_COUNT - 1): - annotation_layers.append( - self.insert_annotation_layer(name=f"name{cx}", descr=f"descr{cx}") - ) - layer_with_annotations = self.insert_annotation_layer( - "layer_with_annotations" - ) - annotation_layers.append(layer_with_annotations) - for cx in range(ANNOTATIONS_COUNT): - annotations.append( - self.insert_annotation( - layer_with_annotations, - short_descr=f"short_descr{cx}", - long_descr=f"long_descr{cx}", - ) - ) - yield annotation_layers - - # rollback changes - for annotation_layer in annotation_layers: - db.session.delete(annotation_layer) - for annotation in annotations: - db.session.delete(annotation) - db.session.commit() - @staticmethod def get_layer_with_annotation() -> AnnotationLayer: return ( @@ -421,9 +359,10 @@ def test_get_annotation(self): """ Annotation API: Test get annotation """ + annotation_id = 1 annotation = ( db.session.query(Annotation) - .filter(Annotation.short_descr == "short_descr1") + .filter(Annotation.short_descr == f"short_descr{annotation_id}") .one_or_none() ) @@ -436,12 +375,12 @@ def test_get_annotation(self): expected_result = { "id": annotation.id, - "end_dttm": None, + "end_dttm": get_end_dttm(annotation_id).isoformat(), "json_metadata": "", "layer": {"id": annotation.layer_id, "name": "layer_with_annotations"}, "long_descr": annotation.long_descr, "short_descr": annotation.short_descr, - "start_dttm": None, + "start_dttm": get_start_dttm(annotation_id).isoformat(), } data = json.loads(rv.data.decode("utf-8")) diff --git a/tests/annotation_layers/fixtures.py b/tests/annotation_layers/fixtures.py new file mode 100644 index 0000000000000..d2960acad5b4d --- /dev/null +++ b/tests/annotation_layers/fixtures.py @@ -0,0 +1,101 @@ +# 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. +# isort:skip_file +import pytest +from datetime import datetime +from typing import Optional + +from superset import db +from superset.models.annotations import Annotation, AnnotationLayer + +from tests.test_app import app + + +ANNOTATION_LAYERS_COUNT = 10 +ANNOTATIONS_COUNT = 5 + + +def get_start_dttm(annotation_id: int) -> datetime: + return datetime(1990 + annotation_id, 1, 1) + + +def get_end_dttm(annotation_id: int) -> datetime: + return datetime(1990 + annotation_id, 7, 1) + + +def _insert_annotation_layer(name: str = "", descr: str = "") -> AnnotationLayer: + annotation_layer = AnnotationLayer(name=name, descr=descr,) + db.session.add(annotation_layer) + db.session.commit() + return annotation_layer + + +def _insert_annotation( + layer: AnnotationLayer, + short_descr: str, + long_descr: str, + json_metadata: Optional[str] = "", + start_dttm: Optional[datetime] = None, + end_dttm: Optional[datetime] = None, +) -> Annotation: + annotation = Annotation( + layer=layer, + short_descr=short_descr, + long_descr=long_descr, + json_metadata=json_metadata, + start_dttm=start_dttm, + end_dttm=end_dttm, + ) + db.session.add(annotation) + db.session.commit() + return annotation + + +@pytest.fixture() +def create_annotation_layers(): + """ + Creates ANNOTATION_LAYERS_COUNT-1 layers with no annotations + and a final one with ANNOTATION_COUNT childs + :return: + """ + with app.app_context(): + annotation_layers = [] + annotations = [] + for cx in range(ANNOTATION_LAYERS_COUNT - 1): + annotation_layers.append( + _insert_annotation_layer(name=f"name{cx}", descr=f"descr{cx}") + ) + layer_with_annotations = _insert_annotation_layer("layer_with_annotations") + annotation_layers.append(layer_with_annotations) + for cx in range(ANNOTATIONS_COUNT): + annotations.append( + _insert_annotation( + layer_with_annotations, + short_descr=f"short_descr{cx}", + long_descr=f"long_descr{cx}", + start_dttm=get_start_dttm(cx), + end_dttm=get_end_dttm(cx), + ) + ) + yield annotation_layers + + # rollback changes + for annotation_layer in annotation_layers: + db.session.delete(annotation_layer) + for annotation in annotations: + db.session.delete(annotation) + db.session.commit() diff --git a/tests/charts/api_tests.py b/tests/charts/api_tests.py index 511186db2d54e..747976943764e 100644 --- a/tests/charts/api_tests.py +++ b/tests/charts/api_tests.py @@ -30,17 +30,18 @@ from sqlalchemy import and_ from sqlalchemy.sql import func -from superset.connectors.sqla.models import SqlaTable -from superset.utils.core import get_example_database -from tests.fixtures.unicode_dashboard import load_unicode_dashboard_with_slice from tests.test_app import app +from superset.connectors.sqla.models import SqlaTable +from superset.utils.core import AnnotationType, get_example_database from superset.connectors.connector_registry import ConnectorRegistry from superset.extensions import db, security_manager +from superset.models.annotations import AnnotationLayer from superset.models.core import Database, FavStar, FavStarClassName from superset.models.dashboard import Dashboard from superset.models.reports import ReportSchedule, ReportScheduleType from superset.models.slice import Slice from superset.utils import core as utils + from tests.base_api_tests import ApiOwnersTestCaseMixin from tests.base_tests import SupersetTestCase from tests.fixtures.importexport import ( @@ -50,7 +51,9 @@ dataset_config, dataset_metadata_config, ) -from tests.fixtures.query_context import get_query_context +from tests.fixtures.query_context import get_query_context, ANNOTATION_LAYERS +from tests.fixtures.unicode_dashboard import load_unicode_dashboard_with_slice +from tests.annotation_layers.fixtures import create_annotation_layers CHART_DATA_URI = "api/v1/chart/data" CHARTS_FIXTURE_COUNT = 10 @@ -1383,3 +1386,44 @@ def test_import_chart_invalid(self): assert response == { "message": {"metadata.yaml": {"type": ["Must be equal to Slice."]}} } + + @pytest.mark.usefixtures("create_annotation_layers") + def test_chart_data_annotations(self): + """ + Chart data API: Test chart data query + """ + self.login(username="admin") + table = self.get_table_by_name("birth_names") + request_payload = get_query_context(table.name, table.id, table.type) + + annotation_layers = [] + request_payload["queries"][0]["annotation_layers"] = annotation_layers + + # formula + annotation_layers.append(ANNOTATION_LAYERS[AnnotationType.FORMULA]) + + # interval + interval_layer = ( + db.session.query(AnnotationLayer) + .filter(AnnotationLayer.name == "name1") + .one() + ) + interval = ANNOTATION_LAYERS[AnnotationType.INTERVAL] + interval["value"] = interval_layer.id + annotation_layers.append(interval) + + # event + event_layer = ( + db.session.query(AnnotationLayer) + .filter(AnnotationLayer.name == "name2") + .one() + ) + event = ANNOTATION_LAYERS[AnnotationType.EVENT] + event["value"] = event_layer.id + annotation_layers.append(event) + + rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data") + self.assertEqual(rv.status_code, 200) + data = json.loads(rv.data.decode("utf-8")) + # response should only contain interval and event data, not formula + self.assertEqual(len(data["result"][0]["annotation_data"]), 2) diff --git a/tests/fixtures/query_context.py b/tests/fixtures/query_context.py index 3881e57641a9d..21bf2b5876e37 100644 --- a/tests/fixtures/query_context.py +++ b/tests/fixtures/query_context.py @@ -17,6 +17,8 @@ import copy from typing import Any, Dict, List +from superset.utils.core import AnnotationType + QUERY_OBJECTS = { "birth_names": { "extras": {"where": "", "time_range_endpoints": ["inclusive", "exclusive"]}, @@ -37,6 +39,77 @@ } } +ANNOTATION_LAYERS = { + AnnotationType.FORMULA: { + "annotationType": "FORMULA", + "color": "#ff7f44", + "hideLine": False, + "name": "my formula", + "opacity": "", + "overrides": {"time_range": None}, + "show": True, + "showMarkers": False, + "sourceType": "", + "style": "solid", + "value": "3+x", + "width": 5, + }, + AnnotationType.EVENT: { + "name": "my event", + "annotationType": "EVENT", + "sourceType": "NATIVE", + "color": "#e04355", + "opacity": "", + "style": "solid", + "width": 5, + "showMarkers": False, + "hideLine": False, + "value": 1, + "overrides": {"time_range": None}, + "show": True, + "titleColumn": "", + "descriptionColumns": [], + "timeColumn": "", + "intervalEndColumn": "", + }, + AnnotationType.INTERVAL: { + "name": "my interval", + "annotationType": "INTERVAL", + "sourceType": "NATIVE", + "color": "#e04355", + "opacity": "", + "style": "solid", + "width": 1, + "showMarkers": False, + "hideLine": False, + "value": 1, + "overrides": {"time_range": None}, + "show": True, + "titleColumn": "", + "descriptionColumns": [], + "timeColumn": "", + "intervalEndColumn": "", + }, + AnnotationType.TIME_SERIES: { + "annotationType": "TIME_SERIES", + "color": None, + "descriptionColumns": [], + "hideLine": False, + "intervalEndColumn": "", + "name": "my line", + "opacity": "", + "overrides": {"time_range": None}, + "show": True, + "showMarkers": False, + "sourceType": "line", + "style": "dashed", + "timeColumn": "", + "titleColumn": "", + "value": 837, + "width": 5, + }, +} + POSTPROCESSING_OPERATIONS = { "birth_names": [ {