#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
import unittest
from copy import deepcopy
from unittest import mock

from airflow.providers.amazon.aws.hooks.glue_crawler import AwsGlueCrawlerHook

mock_crawler_name = 'test-crawler'
mock_role_name = 'test-role'
mock_config = {
    'Name': mock_crawler_name,
    'Description': 'Test glue crawler from Airflow',
    'DatabaseName': 'test_db',
    'Role': mock_role_name,
    'Targets': {
        'S3Targets': [
            {
                'Path': 's3://test-glue-crawler/foo/',
                'Exclusions': [
                    's3://test-glue-crawler/bar/',
                ],
                'ConnectionName': 'test-s3-conn',
            }
        ],
        'JdbcTargets': [
            {
                'ConnectionName': 'test-jdbc-conn',
                'Path': 'test_db/test_table>',
                'Exclusions': [
                    'string',
                ],
            }
        ],
        'MongoDBTargets': [
            {'ConnectionName': 'test-mongo-conn', 'Path': 'test_db/test_collection', 'ScanAll': True}
        ],
        'DynamoDBTargets': [{'Path': 'test_db/test_table', 'scanAll': True, 'scanRate': 123.0}],
        'CatalogTargets': [
            {
                'DatabaseName': 'test_glue_db',
                'Tables': [
                    'test',
                ],
            }
        ],
    },
    'Classifiers': ['test-classifier'],
    'TablePrefix': 'test',
    'SchemaChangePolicy': {
        'UpdateBehavior': 'UPDATE_IN_DATABASE',
        'DeleteBehavior': 'DEPRECATE_IN_DATABASE',
    },
    'RecrawlPolicy': {'RecrawlBehavior': 'CRAWL_EVERYTHING'},
    'LineageConfiguration': 'ENABLE',
    'Configuration': """
    {
        "Version": 1.0,
        "CrawlerOutput": {
            "Partitions": { "AddOrUpdateBehavior": "InheritFromTable" }
        }
    }
    """,
    'SecurityConfiguration': 'test',
    'Tags': {'test': 'foo'},
}


class TestAwsGlueCrawlerHook(unittest.TestCase):
    @classmethod
    def setUp(cls):
        cls.hook = AwsGlueCrawlerHook(aws_conn_id="aws_default")

    def test_init(self):
        self.assertEqual(self.hook.aws_conn_id, "aws_default")

    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    def test_has_crawler(self, mock_get_conn):
        response = self.hook.has_crawler(mock_crawler_name)
        self.assertEqual(response, True)
        mock_get_conn.return_value.get_crawler.assert_called_once_with(Name=mock_crawler_name)

    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    def test_has_crawler_crawled_doesnt_exists(self, mock_get_conn):
        class MockException(Exception):
            pass

        mock_get_conn.return_value.exceptions.EntityNotFoundException = MockException
        mock_get_conn.return_value.get_crawler.side_effect = MockException("AAA")
        response = self.hook.has_crawler(mock_crawler_name)
        self.assertEqual(response, False)
        mock_get_conn.return_value.get_crawler.assert_called_once_with(Name=mock_crawler_name)

    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    def test_update_crawler_needed(self, mock_get_conn):
        mock_get_conn.return_value.get_crawler.return_value = {'Crawler': mock_config}

        mock_config_two = deepcopy(mock_config)
        mock_config_two['Role'] = 'test-2-role'
        response = self.hook.update_crawler(**mock_config_two)
        self.assertEqual(response, True)
        mock_get_conn.return_value.get_crawler.assert_called_once_with(Name=mock_crawler_name)
        mock_get_conn.return_value.update_crawler.assert_called_once_with(**mock_config_two)

    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    def test_update_crawler_not_needed(self, mock_get_conn):
        mock_get_conn.return_value.get_crawler.return_value = {'Crawler': mock_config}
        response = self.hook.update_crawler(**mock_config)
        self.assertEqual(response, False)
        mock_get_conn.return_value.get_crawler.assert_called_once_with(Name=mock_crawler_name)

    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    def test_create_crawler(self, mock_get_conn):
        mock_get_conn.return_value.create_crawler.return_value = {'Crawler': {'Name': mock_crawler_name}}
        glue_crawler = self.hook.create_crawler(**mock_config)
        self.assertIn("Crawler", glue_crawler)
        self.assertIn("Name", glue_crawler["Crawler"])
        self.assertEqual(glue_crawler["Crawler"]["Name"], mock_crawler_name)

    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    def test_start_crawler(self, mock_get_conn):
        result = self.hook.start_crawler(mock_crawler_name)
        self.assertEqual(result, mock_get_conn.return_value.start_crawler.return_value)

        mock_get_conn.return_value.start_crawler.assert_called_once_with(Name=mock_crawler_name)

    @mock.patch.object(AwsGlueCrawlerHook, "get_crawler")
    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    def test_wait_for_crawler_completion_instant_ready(self, mock_get_conn, mock_get_crawler):
        mock_get_crawler.side_effect = [
            {'State': 'READY', 'LastCrawl': {'Status': 'MOCK_STATUS'}},
        ]
        mock_get_conn.return_value.get_crawler_metrics.return_value = {
            'CrawlerMetricsList': [
                {
                    'LastRuntimeSeconds': 'TEST-A',
                    'MedianRuntimeSeconds': 'TEST-B',
                    'TablesCreated': 'TEST-C',
                    'TablesUpdated': 'TEST-D',
                    'TablesDeleted': 'TEST-E',
                }
            ]
        }
        result = self.hook.wait_for_crawler_completion(mock_crawler_name)
        self.assertEqual(result, 'MOCK_STATUS')
        mock_get_conn.assert_has_calls(
            [
                mock.call(),
                mock.call().get_crawler_metrics(CrawlerNameList=[mock_crawler_name]),
            ]
        )
        mock_get_crawler.assert_has_calls(
            [
                mock.call(mock_crawler_name),
            ]
        )

    @mock.patch.object(AwsGlueCrawlerHook, "get_conn")
    @mock.patch.object(AwsGlueCrawlerHook, "get_crawler")
    @mock.patch('airflow.providers.amazon.aws.hooks.glue_crawler.sleep')
    def test_wait_for_crawler_completion_retry_two_times(self, mock_sleep, mock_get_crawler, mock_get_conn):
        mock_get_crawler.side_effect = [
            {'State': 'RUNNING'},
            {'State': 'READY', 'LastCrawl': {'Status': 'MOCK_STATUS'}},
        ]
        mock_get_conn.return_value.get_crawler_metrics.side_effect = [
            {'CrawlerMetricsList': [{'TimeLeftSeconds': 12}]},
            {
                'CrawlerMetricsList': [