# 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, 2, 97, 17), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((59, 42, 60, 63), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((61, 55, 61, 66), 'mock.Mock', 'mock.Mock', ({}, {}), '()', False, 'import mock\n'), ((71, 42, 71, 58), 'mock.MagicMock', 'mock.MagicMock', ({}, {}), '()', False, 'import mock\n'), ((74, 37, 74, 53), 'mock.MagicMock', 'mock.MagicMock', ({}, {}), '()', False, 'import mock\n'), ((63, 25, 64, 44), 'mock.patch.object', 'mock.patch.object', (), '', False, 'import mock\n'), ((85, 15, 86, 38), 'dags.bq_to_cm_dag.BigQueryToCMDag', 'bq_to_cm_dag.BigQueryToCMDag', ({(86, 8, 86, 37): "AIRFLOW_VARIABLES['dag_name']"}, {}), "(AIRFLOW_VARIABLES['dag_name'])", False, 'from dags import bq_to_cm_dag\n')]