{# Copyright 2021 Curtin University
#
# Licensed 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.
# Author: Tuan Chien, Aniek Roelofs -#}
import os
import pendulum

from {{ package_name }}.workflows.{{ workflow_module }} import {{ workflow_class }}
from observatory.platform.utils.config_utils import module_file_path
from observatory.platform.utils.test_utils import (
    ObservatoryEnvironment,
    ObservatoryTestCase,
)


class Test{{ workflow_class }}(ObservatoryTestCase):
    """ Tests for the {{ workflow_class }} workflow """

    def __init__(self, *args, **kwargs):
        """Constructor which sets up variables used by tests.

        :param args: arguments.
        :param kwargs: keyword arguments.
        """
        super(Test{{ workflow_class }}, self).__init__(*args, **kwargs)
        self.project_id = os.getenv("TEST_GCP_PROJECT_ID")
        self.data_location = os.getenv("TEST_GCP_DATA_LOCATION")

        self.execution_date = pendulum.datetime(year=2021, month=1, day=1)

    def test_dag_structure(self):
        """Test that the DAG has the correct structure.

        :return: None
        """

        dag = {{ workflow_class }}().make_dag()
        self.assert_dag_structure({"check_dependencies": ["task1"], "task1": ["cleanup"], "cleanup": []}, dag)

    def test_dag_load(self):
        """Test that the DAG can be loaded from a DAG bag.

        :return: None
        """

        with ObservatoryEnvironment().create():
            dag_file = os.path.join(module_file_path("{{ package_name }}.dags"), "{{ workflow_module }}.py")
            self.assert_dag_load("{{ workflow_module }}", dag_file)

    def test_workflow(self):
        """Test the workflow end to end.

        :return: None.
        """
        # Setup Observatory environment
        env = ObservatoryEnvironment(self.project_id, self.data_location)
        dataset_id = env.add_dataset()

        # Setup Workflow
        workflow = {{ workflow_class }}()
        dag = workflow.make_dag()

        # Create the Observatory environment and run tests
        with env.create(task_logging=True):
            with env.create_dag_run(dag, self.execution_date):
                env.run_task(workflow.check_dependencies.__name__)
                env.run_task(workflow.task1.__name__)
                env.run_task(workflow.cleanup.__name__)
