# python3
# coding=utf-8
# Copyright 2020 Google LLC.
#
# 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.

"""Tests for dags.bq_to_cm_dag."""

import unittest

from airflow.contrib.hooks import bigquery_hook
from airflow.models import baseoperator
from airflow.models import dag
from airflow.models import variable
import mock

from gps_building_blocks.cloud.utils import cloud_auth
from dags import bq_to_cm_dag
from plugins.pipeline_plugins.hooks import monitoring_hook

_DAG_NAME = bq_to_cm_dag._DAG_NAME

AIRFLOW_VARIABLES = {
    'dag_name': _DAG_NAME,
    f'{_DAG_NAME}_schedule': '@once',
    f'{_DAG_NAME}_retries': 0,
    f'{_DAG_NAME}_retry_delay': 3,
    f'{_DAG_NAME}_is_retry': True,
    f'{_DAG_NAME}_is_run': True,
    f'{_DAG_NAME}_enable_run_report': False,
    f'{_DAG_NAME}_enable_monitoring': True,
    f'{_DAG_NAME}_enable_monitoring_cleanup': False,
    'monitoring_data_days_to_live': 50,
    'monitoring_dataset': 'test_monitoring_dataset',
    'monitoring_table': 'test_monitoring_table',
    'monitoring_bq_conn_id': 'test_monitoring_conn',
    'bq_dataset_id': 'test_dataset',
    'bq_table_id': 'test_table',
    'cm_profile_id': 'cm_profile_id',
    'cm_service_account': 'cm_service_account'
}


class BQToCMDAGTest(unittest.TestCase):

  def setUp(self):
    super(BQToCMDAGTest, self).setUp()
    self.addCleanup(mock.patch.stopall)
    self.build_impersonated_client_mock = mock.patch.object(
        cloud_auth, 'build_impersonated_client', autospec=True)
    self.build_impersonated_client_mock.return_value = mock.Mock()
    self.build_impersonated_client_mock.start()
    self.mock_variable = mock.patch.object(
        variable, 'Variable', autospec=True).start()
    # `side_effect` is assigned to `lambda` to dynamically return values
    # each time when self.mock_variable is called.
    self.mock_variable.get.side_effect = (
        lambda key, value: AIRFLOW_VARIABLES[key])

    self.original_bigquery_hook_init = bigquery_hook.BigQueryHook.__init__
    bigquery_hook.BigQueryHook.__init__ = mock.MagicMock()

    self.original_monitoring_hook = monitoring_hook.MonitoringHook
    monitoring_hook.MonitoringHook = mock.MagicMock()

  def tearDown(self):
    super().tearDown()
    bigquery_hook.BigQueryHook.__init__ = self.original_bigquery_hook_init
    monitoring_hook.MonitoringHook = self.original_monitoring_hook

  def test_create_dag(self):
    """Tests that returned DAG contains correct DAG and tasks."""
    expected_task_ids = ['bq_to_cm_retry_task', 'bq_to_cm_task']

    test_dag = bq_to_cm_dag.BigQueryToCMDag(
        AIRFLOW_VARIABLES['dag_name']).create_dag()

    self.assertIsInstance(test_dag, dag.DAG)
    self.assertEqual(len(test_dag.tasks), len(expected_task_ids))
    for task in test_dag.tasks:
      self.assertIsInstance(task, baseoperator.BaseOperator)
    actual_task_ids = [t.task_id for t in test_dag.tasks]
    self.assertListEqual(actual_task_ids, expected_task_ids)


if __name__ == '__main__':
  unittest.main()

[(97, 'unittest.main', 'unittest.main', False, 'import unittest\n'), (59, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (61, 'mock.Mock', 'mock.Mock', False, 'import mock\n'), (71, 'mock.MagicMock', 'mock.MagicMock', False, 'import mock\n'), (74, 'mock.MagicMock', 'mock.MagicMock', False, 'import mock\n'), (63, 'mock.patch.object', 'mock.patch.object', False, 'import mock\n'), (85, 'dags.bq_to_cm_dag.BigQueryToCMDag', 'bq_to_cm_dag.BigQueryToCMDag', False, 'from dags import bq_to_cm_dag\n')]