tests/annotation_layers/fixtures.py

# 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()