{# 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 unittest.mock import patch
from airflow.models import Connection

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,
)
from observatory.api.client import ApiClient, Configuration
from observatory.api.client.api.observatory_api import ObservatoryApi  # noqa: E501
from observatory.api.testing import ObservatoryApiEnvironment
from observatory.api.client.model.organisation import Organisation
from observatory.api.client.model.table_type import TableType
from observatory.api.client.model.dataset_type import DatasetType
from observatory.api.client.model.workflow import Workflow
from observatory.api.client.model.workflow_type import WorkflowType
from observatory.api.client.model.dataset import Dataset


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.first_execution_date = pendulum.datetime(year=2021, month=1, day=1)
        self.second_execution_date = pendulum.datetime(year=2021, month=2, day=1)

        # API environment
        self.host = "localhost"
        self.port = 5001
        configuration = Configuration(host=f"http://{self.host}:{self.port}")
        api_client = ApiClient(configuration)
        self.api = ObservatoryApi(api_client=api_client)  # noqa: E501
        self.env = ObservatoryApiEnvironment(host=self.host, port=self.port)
        self.workflow_id = 1
        self.org_name = "Curtin University"

    def setup_api(self):
        org = Organisation(name=self.org_name)
        result = self.api.put_organisation(org)
        self.assertIsInstance(result, Organisation)

        tele_type = WorkflowType(type_id="my_stream", name="My Workflow")
        result = self.api.put_workflow_type(tele_type)
        self.assertIsInstance(result, WorkflowType)

        workflow = Workflow(organisation=Organisation(id=1), workflow_type=WorkflowType(id=1))
        result = self.api.put_workflow(workflow)
        self.assertIsInstance(result, Workflow)

        table_type = TableType(
                type_id="partitioned",
                name="partitioned bq table",
        )
        self.api.put_table_type(table_type)

        dataset_type = DatasetType(
            name="My dataset type",
            type_id="type id",
            table_type=TableType(id=1),
        )
        
        self.api.put_dataset_type(dataset_type)

        dataset = Dataset(name="My dataset", service="bigquery", address="project.dataset.table", workflow=Workflow(id=1), dataset_type=DatasetType(id=1))
        result = self.api.put_dataset(dataset)
        self.assertIsInstance(result, Dataset)

    def setup_connections(self, env):
        """Setup the Observatory API airflow connection.

        :return: None
        """
        # Add Observatory API connection
        conn = Connection(conn_id=AirflowConns.OBSERVATORY_API, uri=f"http://:password@{self.host}:{self.port}")
        env.add_connection(conn)

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

        :return: None
        """
        dag = {{ workflow_class }}(workflow_id=self.workflow_id).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, api_host=self.host, api_port=self.port)

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

    @patch("observatory.platform.utils.release_utils.make_observatory_api")
    def test_workflow(self, m_makeapi):
        """Test the workflow end to end.

        :return: None.
        """

        m_makeapi.return_value = self.api

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

        # Setup Workflow
        workflow = {{ workflow_class }}(dataset_id=dataset_id, workflow_id=self.workflow_id)
        dag = workflow.make_dag()

        # Test the workflow tasks
        with env.create(task_logging=True):
            self.setup_api()

            # Test the dag run for the first release
            with env.create_dag_run(dag, self.first_execution_date):
                env.run_task(workflow.check_dependencies.__name__)
                env.run_task(workflow.task1.__name__)
                env.run_task(workflow.cleanup.__name__)

            # Test the dag run for a later release
            with env.create_dag_run(dag, self.second_execution_date):
                env.run_task(workflow.check_dependencies.__name__)
                env.run_task(workflow.task1.__name__)
                env.run_task(workflow.cleanup.__name__)
