# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
from unittest.mock import ANY, MagicMock, Mock

import pytest

from czc import s3
from czt.unittest import Context, MockEx


def _list_response(keys=[], has_more=False, page_token=None, prefix_keys=[]):
    return {
        'Contents': [{'Key': k} for k in keys],
        'IsTruncated': has_more,
        'NextContinuationToken': page_token,
        'CommonPrefixes': [{'Prefix': k} for k in prefix_keys],
    }


@pytest.fixture(scope='function')
def context(mocker):
    mock_aws_service = mocker.patch.object(s3, 'aws_service', autospec=False)
    extra = {'paginator': Mock(), 'mock_aws_service': mock_aws_service}
    with Context(mocker, s3, [], extra) as context:
        context.paginator.paginate.return_value = [_list_response()]
        context.mock_aws_service.s3_client.get_paginator.return_value = context.paginator
        yield context


@pytest.mark.unit
@pytest.mark.parametrize(
    'bucket,key',
    [
        ('some-bucket', 'abc'),
        ('other-bucket', 'xyz'),
    ],
)
def test_get_presigned_url_should_return_link_for_bucket_and_key(context, bucket, key):
    result = s3.get_presigned_url(bucket, key)
    assert result == context.mock_aws_service.s3_client.generate_presigned_url.return_value
    assert MockEx(context.mock_aws_service.s3_client.generate_presigned_url).was_called_once_with(
        'get_object', Params={'Bucket': bucket, 'Key': key}, ExpiresIn=900, HttpMethod='GET'
    )


@pytest.mark.unit
@pytest.mark.parametrize('expires_in', [42, 300])
def test_get_presigned_url_should_set_requested_expiration(context, expires_in):
    result = s3.get_presigned_url('bucket', 'key', expires_in=expires_in)
    assert result == context.mock_aws_service.s3_client.generate_presigned_url.return_value
    assert MockEx(context.mock_aws_service.s3_client.generate_presigned_url).was_called_once_with(
        'get_object', Params=ANY, ExpiresIn=expires_in, HttpMethod='GET'
    )


@pytest.mark.unit
@pytest.mark.parametrize('bucket, key, data, expected_object', [('some-bucket', 'abc', '{"foo": "bar"}', {'foo': 'bar'})])
def test_get_json_object_should_deserialize_object_from_s3(context, bucket, key, data, expected_object):
    mock_reader = MagicMock()
    mock_reader.read.return_value = data

    MockEx(context.mock_aws_service.s3_client.get_object).if_called_with(Bucket=bucket, Key=key).return_value(
        {'Body': mock_reader}
    )
    result = s3.get_json_object(bucket, key)
    assert result == expected_object


@pytest.mark.unit
@pytest.mark.parametrize('bucket, key, obj, expected_data', [('some-bucket', 'abc', {'foo': 'bar'}, b'{"foo": "bar"}')])
def test_put_json_object_should_serialize_object_to_s3(context, bucket, key, obj, expected_data):
    s3.put_json_object(bucket, key, obj)
    assert MockEx(context.mock_aws_service.s3_client.put_object).was_called_once_with(
        Bucket=bucket, Key=key, Body=expected_data
    )


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize('bucket,prefix', [('my-bucket', 'foo/bar'), ('another-bucket', 'stuff/goes/here')])
def test_list_objects_should_request_expected_bucket_and_prefix(context, bucket, prefix):
    s3.list_objects(bucket, prefix)
    assert MockEx(context.paginator.paginate).was_called_once_with(Bucket=bucket, Prefix=prefix)


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize(
    's3_responses,expected_result',
    [
        ([_list_response(keys=['some_key', 'another_key'])], ['some_key', 'another_key']),
        ([_list_response(keys=['a', 'b'], has_more=True), _list_response(keys=['c', 'd'])], ['a', 'b', 'c', 'd']),
        ([_list_response(has_more=True), _list_response(keys=['key', 'key', 'again'])], ['key', 'key', 'again']),
    ],
)
def test_list_objects_should_flatten_s3_response_contents(context, s3_responses, expected_result):
    context.paginator.paginate.return_value = s3_responses

    result = s3.list_objects('bucket', 'prefix')
    assert result == expected_result


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize(
    's3_responses,expected_result',
    [
        ([_list_response(prefix_keys=['foo/bar', 'foo/baz'])], ['foo/bar', 'foo/baz']),
        ([_list_response(prefix_keys=['some/key']), _list_response(prefix_keys=['other/key'])], ['some/key', 'other/key']),
    ],
)
def test_list_common_prefixes(context, s3_responses, expected_result):
    context.paginator.paginate.return_value = s3_responses

    result = s3.list_common_prefixes('bucket', 'prefix', 'delimiter')
    assert result == expected_result
