# Copyright 2014 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.

from argparse import Namespace

from botocore.exceptions import ClientError

import awscli
from awscli.compat import ZIP_COMPRESSION_MODE, StringIO
from awscli.customizations.codedeploy.push import Push
from awscli.customizations.exceptions import ParamValidationError
from awscli.testutils import mock, unittest


class TestPush(unittest.TestCase):
    def setUp(self):
        self.application_name = 'MyApp'
        self.description = 'MyApp revision'
        self.source = '/tmp'
        self.appspec = 'appspec.yml'
        self.appspec_path = f'{self.source}/{self.appspec}'
        self.bucket = 'foo'
        self.key = 'bar/baz.zip'
        self.s3_location = 's3://' + self.bucket + '/' + self.key
        self.eTag = '"1a2b3cd45e"'
        self.version_id = '12341234-1234-1234-1234-123412341234'
        self.upload_id = 'upload_id'
        self.region = 'us-east-1'
        self.endpoint_url = 'https://codedeploy.aws.amazon.com'

        self.args = Namespace()
        self.args.application_name = self.application_name
        self.args.s3_location = self.s3_location
        self.args.ignore_hidden_files = False
        self.args.no_ignore_hidden_files = False
        self.args.description = self.description
        self.args.source = self.source

        self.globals = Namespace()
        self.globals.region = self.region
        self.globals.endpoint_url = self.endpoint_url
        self.globals.verify_ssl = False

        self.upload_response = {
            'ETag': self.eTag,
            'VersionId': self.version_id,
        }
        self.revision = {
            'revisionType': 'S3',
            's3Location': {
                'bucket': self.bucket,
                'key': self.key,
                'bundleType': 'zip',
                'eTag': self.eTag,
                'version': self.version_id,
            },
        }

        self.bundle_mock = mock.MagicMock()
        self.bundle_mock.tell.return_value = 5 << 20
        self.bundle_mock.read.return_value = b'a' * (5 << 20)
        self.bundle_mock.__enter__.return_value = self.bundle_mock
        self.bundle_mock.__exit__.return_value = None

        self.zipfile_mock = mock.MagicMock()
        self.zipfile_mock.write.return_value = None
        self.zipfile_mock.close.return_value = None
        self.zipfile_mock.__enter__.return_value = self.zipfile_mock
        self.zipfile_mock.__exit__.return_value = None

        self.session = mock.MagicMock()

        self.push = Push(self.session)
        self.push.s3 = mock.MagicMock()
        self.push.s3.put_object.return_value = self.upload_response
        self.push.s3.create_multipart_upload.return_value = {
            'UploadId': self.upload_id
        }
        self.push.s3.upload_part.return_value = {'ETag': self.eTag}
        self.push.s3.complete_multipart_upload.return_value = (
            self.upload_response
        )
        self.push.codedeploy = mock.MagicMock()

    def test_run_main_throws_on_invalid_args(self):
        self.push._validate_args = mock.MagicMock()
        self.push._validate_args.side_effect = RuntimeError()
        with self.assertRaises(RuntimeError):
            self.push._run_main(self.args, self.globals)

    def test_run_main_creates_clients(self):
        self.push._validate_args = mock.MagicMock()
        self.push._push = mock.MagicMock()
        self.push._run_main(self.args, self.globals)
        self.session.create_client.assert_has_calls(
            [
                mock.call(
                    'codedeploy',
                    region_name=self.region,
                    endpoint_url=self.endpoint_url,
                    verify=self.globals.verify_ssl,
                ),
                mock.call('s3', region_name=self.region),
            ]
        )

    def test_run_main_calls_push(self):
        self.push._validate_args = mock.MagicMock()
        self.push._push = mock.MagicMock()
        self.push._run_main(self.args, self.globals)
        self.push._push.assert_called_with(self.args)

    @mock.patch.object(
        awscli.customizations.codedeploy.push, 'validate_s3_location'
    )
    def test_validate_args_throws_on_invalid_s3_url(
        self, validate_s3_location
    ):
        self.args.s3_location = 's3:/foo/bar/baz'
        validate_s3_location.side_effect = RuntimeError()
        with self.assertRaises(RuntimeError):
            self.push._validate_args(self.args)

    def test_validate_args_throws_on_ignore_and_no_ignore_hidden_files(self):
        self.args.ignore_hidden_files = True
        self.args.no_ignore_hidden_files = True
        with self.assertRaises(ParamValidationError):
            self.push._validate_args(self.args)

    def test_validate_args_default_description(self):
        self.args.description = None
        self.push._validate_args(self.args)
        self.assertRegex(self.args.description, 'Uploaded by AWS CLI .* UTC')

    def test_push_throws_on_upload_to_s3_error(self):
        self.args.bucket = self.bucket
        self.args.key = self.key
        self.push._compress = mock.MagicMock(return_value=self.bundle_mock)
        self.push._upload_to_s3 = mock.MagicMock()
        self.push._upload_to_s3.side_effect = RuntimeError()
        with self.assertRaises(RuntimeError):
            self.push._push(self.args)

    def test_push_strips_quotes_from_etag(self):
        self.args.bucket = self.bucket
        self.args.key = self.key
        self.push._compress = mock.MagicMock(return_value=self.bundle_mock)
        self.push._upload_to_s3 = mock.MagicMock(
            return_value=self.upload_response
        )
        self.push._register_revision = mock.MagicMock()
        self.push._push(self.args)
        self.push._register_revision.assert_called_with(self.args)
        self.assertEqual(
            str(self.args.eTag), self.upload_response['ETag'].replace('"', "")
        )

    @mock.patch('sys.stdout', new_callable=StringIO)
    def test_push_output_message(self, stdout_mock):
        self.args.bucket = self.bucket
        self.args.key = self.key
        self.push._compress = mock.MagicMock(return_value=self.bundle_mock)
        self.push._upload_to_s3 = mock.MagicMock(
            return_value=self.upload_response
        )
        self.push._register_revision = mock.MagicMock()
        self.push._push(self.args)
        output = stdout_mock.getvalue().strip()
        expected_revision_output = (
            '--s3-location bucket={0},key={1},'
            'bundleType=zip,eTag={2},version={3}'.format(
                self.bucket,
                self.key,
                self.eTag.replace('"', ""),
                self.version_id,
            )
        )
        expected_output = (
            'To deploy with this revision, run:\n'
            'aws deploy create-deployment '
            f'--application-name {self.application_name} {expected_revision_output} '
            '--deployment-group-name <deployment-group-name> '
            '--deployment-config-name <deployment-config-name> '
            '--description <description>'
        )
        self.assertEqual(expected_output, output)

    @mock.patch('zipfile.ZipFile')
    @mock.patch('tempfile.TemporaryFile')
    @mock.patch('os.path')
    @mock.patch('os.walk')
    def test_compress_throws_when_no_appspec(self, walk, path, tf, zf):
        walk.return_value = [(self.source, [], ['noappspec.yml'])]
        noappsec_path = self.source + '/noappspec.yml'
        path.join.return_value = noappsec_path
        path.sep = '/'
        path.abspath.side_effect = [self.source, noappsec_path]
        tf.return_value = self.bundle_mock
        zf.return_value = self.zipfile_mock
        with self.assertRaises(RuntimeError):
            with self.push._compress(
                self.args.source, self.args.ignore_hidden_files
            ):
                pass

    @mock.patch('zipfile.ZipFile')
    @mock.patch('tempfile.TemporaryFile')
    @mock.patch('os.path')
    @mock.patch('os.walk')
    def test_compress_writes_to_zip_file(self, walk, path, tf, zf):
        walk.return_value = [(self.source, [], [self.appspec])]
        path.join.return_value = self.appspec_path
        path.sep = '/'
        path.abspath.side_effect = [self.source, self.appspec_path]
        tf.return_value = self.bundle_mock
        zf.return_value = self.zipfile_mock
        with self.push._compress(
            self.args.source, self.args.ignore_hidden_files
        ):
            zf.assert_called_with(mock.ANY, 'w', allowZip64=True)
            zf().write.assert_called_with(
                '/tmp/appspec.yml', self.appspec, ZIP_COMPRESSION_MODE
            )

    def test_upload_to_s3_with_put_object(self):
        self.args.bucket = self.bucket
        self.args.key = self.key
        response = self.push._upload_to_s3(self.args, self.bundle_mock)
        self.assertDictEqual(self.upload_response, response)
        self.push.s3.put_object.assert_called_with(
            Bucket=self.bucket, Key=self.key, Body=self.bundle_mock
        )
        self.assertFalse(self.push.s3.create_multipart_upload.called)
        self.assertFalse(self.push.s3.upload_part.called)
        self.assertFalse(self.push.s3.complete_multipart_upload.called)
        self.assertFalse(self.push.s3.abort_multipart_upload.called)

    def test_upload_to_s3_with_multipart_upload(self):
        self.args.bucket = self.bucket
        self.args.key = self.key
        self.bundle_mock.tell.return_value = 6 << 20
        self.bundle_mock.read.return_value = b'a' * (6 << 20)
        response = self.push._upload_to_s3(self.args, self.bundle_mock)
        self.assertDictEqual(self.upload_response, response)
        self.assertFalse(self.push.s3.put_object.called)
        self.push.s3.create_multipart_upload.assert_called_with(
            Bucket=self.bucket, Key=self.key
        )
        self.push.s3.upload_part.assert_called_with(
            Bucket=self.bucket,
            Key=self.key,
            UploadId=self.upload_id,
            PartNumber=1,
            Body=mock.ANY,
        )
        self.push.s3.complete_multipart_upload.assert_called_with(
            Bucket=self.bucket,
            Key=self.key,
            UploadId=self.upload_id,
            MultipartUpload={'Parts': [{'PartNumber': 1, 'ETag': self.eTag}]},
        )
        self.assertFalse(self.push.s3.abort_multipart_upload.called)

    def test_upload_to_s3_with_multipart_upload_aborted_on_error(self):
        self.args.bucket = self.bucket
        self.args.key = self.key
        self.bundle_mock.tell.return_value = 6 << 20
        self.bundle_mock.read.return_value = b'a' * (6 << 20)
        self.push.s3.upload_part.side_effect = ClientError(
            {'Error': {'Code': 'Error', 'Message': 'Error'}}, 'UploadPart'
        )
        with self.assertRaises(ClientError):
            self.push._upload_to_s3(self.args, self.bundle_mock)
        self.assertFalse(self.push.s3.put_object.called)
        self.push.s3.create_multipart_upload.assert_called_with(
            Bucket=self.bucket, Key=self.key
        )
        self.assertTrue(self.push.s3.upload_part.called)
        self.assertFalse(self.push.s3.complete_multipart_upload.called)
        self.push.s3.abort_multipart_upload.assert_called_with(
            Bucket=self.bucket, Key=self.key, UploadId=self.upload_id
        )

    def test_register_revision(self):
        self.args.bucket = self.bucket
        self.args.key = self.key
        self.args.eTag = self.eTag
        self.args.version = self.version_id
        self.push._register_revision(self.args)
        self.push.codedeploy.register_application_revision.assert_called_with(
            applicationName=self.application_name,
            description=self.description,
            revision=self.revision,
        )


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