# Copyright 2015 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
#
# https://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.
import io

import pytest
from botocore.exceptions import ClientError

from boto3.s3 import inject
from boto3.s3.transfer import TransferConfig
from tests import mock, unittest


class TestInjectTransferMethods(unittest.TestCase):
    def test_inject_upload_download_file_to_client(self):
        class_attributes = {}
        inject.inject_s3_transfer_methods(class_attributes=class_attributes)
        assert 'upload_file' in class_attributes
        assert 'download_file' in class_attributes

    def test_upload_file_proxies_to_transfer_object(self):
        with mock.patch('boto3.s3.inject.S3Transfer') as transfer:
            inject.upload_file(
                mock.sentinel.CLIENT,
                Filename='filename',
                Bucket='bucket',
                Key='key',
            )
            transfer_in_context_manager = (
                transfer.return_value.__enter__.return_value
            )
            transfer_in_context_manager.upload_file.assert_called_with(
                filename='filename',
                bucket='bucket',
                key='key',
                extra_args=None,
                callback=None,
            )

    def test_download_file_proxies_to_transfer_object(self):
        with mock.patch('boto3.s3.inject.S3Transfer') as transfer:
            inject.download_file(
                mock.sentinel.CLIENT,
                Bucket='bucket',
                Key='key',
                Filename='filename',
            )
            transfer_in_context_manager = (
                transfer.return_value.__enter__.return_value
            )
            transfer_in_context_manager.download_file.assert_called_with(
                bucket='bucket',
                key='key',
                filename='filename',
                extra_args=None,
                callback=None,
            )


class TestBucketLoad(unittest.TestCase):
    def setUp(self):
        self.client = mock.Mock()
        self.resource = mock.Mock()
        self.resource.meta.client = self.client

    def test_bucket_load_finds_bucket(self):
        self.resource.name = 'MyBucket'
        self.client.list_buckets.return_value = {
            'Buckets': [
                {'Name': 'NotMyBucket', 'CreationDate': 1},
                {'Name': self.resource.name, 'CreationDate': 2},
            ],
        }

        inject.bucket_load(self.resource)
        assert self.resource.meta.data == {
            'Name': self.resource.name,
            'CreationDate': 2,
        }

    def test_bucket_load_doesnt_find_bucket(self):
        self.resource.name = 'MyBucket'
        self.client.list_buckets.return_value = {
            'Buckets': [
                {'Name': 'NotMyBucket', 'CreationDate': 1},
                {'Name': 'NotMine2', 'CreationDate': 2},
            ],
        }
        inject.bucket_load(self.resource)
        assert self.resource.meta.data == {}

    def test_bucket_load_encounters_access_exception(self):
        self.client.list_buckets.side_effect = ClientError(
            {'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}},
            'ListBuckets',
        )
        inject.bucket_load(self.resource)
        assert self.resource.meta.data == {}

    def test_bucket_load_encounters_other_exception(self):
        self.client.list_buckets.side_effect = ClientError(
            {
                'Error': {
                    'Code': 'ExpiredToken',
                    'Message': 'The provided token has expired.',
                }
            },
            'ListBuckets',
        )
        with pytest.raises(ClientError):
            inject.bucket_load(self.resource)


class TestBucketTransferMethods(unittest.TestCase):
    def setUp(self):
        self.bucket = mock.Mock(name='my_bucket')
        self.copy_source = {'Bucket': 'foo', 'Key': 'bar'}

    def test_upload_file_proxies_to_meta_client(self):
        inject.bucket_upload_file(self.bucket, Filename='foo', Key='key')
        self.bucket.meta.client.upload_file.assert_called_with(
            Filename='foo',
            Bucket=self.bucket.name,
            Key='key',
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )

    def test_download_file_proxies_to_meta_client(self):
        inject.bucket_download_file(self.bucket, Key='key', Filename='foo')
        self.bucket.meta.client.download_file.assert_called_with(
            Bucket=self.bucket.name,
            Key='key',
            Filename='foo',
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )

    def test_copy(self):
        inject.bucket_copy(self.bucket, self.copy_source, Key='key')
        self.bucket.meta.client.copy.assert_called_with(
            CopySource=self.copy_source,
            Bucket=self.bucket.name,
            Key='key',
            ExtraArgs=None,
            Callback=None,
            SourceClient=None,
            Config=None,
        )

    def test_upload_fileobj(self):
        fileobj = io.BytesIO(b'foo')
        inject.bucket_upload_fileobj(self.bucket, Key='key', Fileobj=fileobj)
        self.bucket.meta.client.upload_fileobj.assert_called_with(
            Bucket=self.bucket.name,
            Fileobj=fileobj,
            Key='key',
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )

    def test_download_fileobj(self):
        obj = io.BytesIO()
        inject.bucket_download_fileobj(self.bucket, Key='key', Fileobj=obj)
        self.bucket.meta.client.download_fileobj.assert_called_with(
            Bucket=self.bucket.name,
            Key='key',
            Fileobj=obj,
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )


class TestObjectTransferMethods(unittest.TestCase):
    def setUp(self):
        self.obj = mock.Mock(bucket_name='my_bucket', key='my_key')
        self.copy_source = {'Bucket': 'foo', 'Key': 'bar'}

    def test_upload_file_proxies_to_meta_client(self):
        inject.object_upload_file(self.obj, Filename='foo')
        self.obj.meta.client.upload_file.assert_called_with(
            Filename='foo',
            Bucket=self.obj.bucket_name,
            Key=self.obj.key,
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )

    def test_download_file_proxies_to_meta_client(self):
        inject.object_download_file(self.obj, Filename='foo')
        self.obj.meta.client.download_file.assert_called_with(
            Bucket=self.obj.bucket_name,
            Key=self.obj.key,
            Filename='foo',
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )

    def test_copy(self):
        inject.object_copy(self.obj, self.copy_source)
        self.obj.meta.client.copy.assert_called_with(
            CopySource=self.copy_source,
            Bucket=self.obj.bucket_name,
            Key=self.obj.key,
            ExtraArgs=None,
            Callback=None,
            SourceClient=None,
            Config=None,
        )

    def test_upload_fileobj(self):
        fileobj = io.BytesIO(b'foo')
        inject.object_upload_fileobj(self.obj, Fileobj=fileobj)
        self.obj.meta.client.upload_fileobj.assert_called_with(
            Bucket=self.obj.bucket_name,
            Fileobj=fileobj,
            Key=self.obj.key,
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )

    def test_download_fileobj(self):
        fileobj = io.BytesIO()
        inject.object_download_fileobj(self.obj, Fileobj=fileobj)
        self.obj.meta.client.download_fileobj.assert_called_with(
            Bucket=self.obj.bucket_name,
            Key=self.obj.key,
            Fileobj=fileobj,
            ExtraArgs=None,
            Callback=None,
            Config=None,
        )


class TestObejctSummaryLoad(unittest.TestCase):
    def setUp(self):
        self.client = mock.Mock()
        self.resource = mock.Mock()
        self.resource.meta.client = self.client
        self.head_object_response = {'ContentLength': 5, 'ETag': 'my-etag'}
        self.client.head_object.return_value = self.head_object_response

    def test_object_summary_load(self):
        inject.object_summary_load(self.resource)
        assert self.resource.meta.data == {'Size': 5, 'ETag': 'my-etag'}

    def test_can_handle_missing_content_length(self):
        self.head_object_response.pop('ContentLength')
        inject.object_summary_load(self.resource)
        assert self.resource.meta.data == {'ETag': 'my-etag'}


def test_disable_threading_if_append_mode(caplog, tmp_path):
    config = TransferConfig(use_threads=True)
    with open(tmp_path / 'myfile', 'ab') as f:
        inject.disable_threading_if_append_mode(config, f)
    assert config.use_threads is False
    assert 'A single thread will be used' in caplog.text


def test_threading_not_disabled_if_not_append_mode(caplog, tmp_path):
    config = TransferConfig(use_threads=True)
    with open(tmp_path / 'myfile', 'wb') as f:
        inject.disable_threading_if_append_mode(config, f)
    assert config.use_threads is True
    assert 'A single thread will be used' not in caplog.text
