{# 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 airflow.models.connection import Connection

from observatory.api.client.model.organisation import Organisation
from observatory.api.server import orm
from {{ package_name }}.identifiers import WorkflowTypes
from {{ package_name }}.workflows.{{ workflow_module }} import {{ workflow_class }}
from observatory.platform.utils.airflow_utils import AirflowConns
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.host = "localhost"
        self.api_port = 5000

        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
        """
        organisation = Organisation(name="Organisation Name")
        dag = {{ workflow_class }}(organisation, "value1").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
        """
        env = ObservatoryEnvironment(self.project_id, self.data_location)
        with env.create():
            # Add Observatory API connection
            conn = Connection(conn_id=AirflowConns.OBSERVATORY_API, uri=f"http://:password@{self.host}:{self.api_port}")
            env.add_connection(conn)

            # Add a workflow
            dt = pendulum.now("UTC")
            workflow_type = orm.WorkflowType(
                name="{{ workflow_class}} Workflow", type_id=WorkflowTypes.{{ workflow_module }}, created=dt, modified=dt
            )
            env.api_session.add(workflow_type)
            organisation = orm.Organisation(name="Organisation", created=dt, modified=dt)
            env.api_session.add(organisation)
            workflow = orm.Workflow(
                name="Organisation {{ workflow_class }} Workflow",
                workflow_type=workflow_type,
                organisation=organisation,
                modified=dt,
                created=dt,
                extra={"key1": "value1"},
            )
            env.api_session.add(workflow)
            env.api_session.commit()

            dag_file = os.path.join(module_file_path("{{ package_name }}.dags"), "{{ workflow_module }}.py")
            self.assert_dag_load("{{ workflow_module }}_organisation", dag_file)

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

        :return: None.
        """
        # Set up organisation name and workflow extra values
        self.organisation_name = "Organisation"
        self.extra = {"extra1": "value1"}
        self.extra1 = self.extra.get("extra1")

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

        # Setup Workflow
        organisation = Organisation(
            name=self.organisation_name,
            project_id=self.project_id,
            download_bucket=env.download_bucket,
            transform_bucket=env.transform_bucket,
        )
        workflow = {{ workflow_class }}(organisation=organisation, extra1=self.extra1, dataset_id=dataset_id)
        dag = workflow.make_dag()

        # Test the workflow tasks
        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__)
