import pytest

import botocore.session
from botocore import UNSIGNED
from botocore.client import Config
from botocore.exceptions import ClientError
from pynamodb.attributes import UnicodeAttribute
from pynamodb.models import Model

from aws_xray_sdk.core import patch
from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core.context import Context

patch(('pynamodb',))


@pytest.fixture(autouse=True)
def construct_ctx():
    """
    Clean up context storage on each test run and begin a segment
    so that later subsegment can be attached. After each test run
    it cleans up context storage again.
    """
    xray_recorder.configure(service='test', sampling=False, context=Context())
    xray_recorder.clear_trace_entities()
    xray_recorder.begin_segment('name')
    yield
    xray_recorder.clear_trace_entities()


def test_exception():
    class SampleModel(Model):
        class Meta:
            region = 'us-west-2'
            table_name = 'mytable'

        sample_attribute = UnicodeAttribute(hash_key=True)

    try:
        SampleModel.describe_table()
    except Exception:
        pass

    subsegments = xray_recorder.current_segment().subsegments
    assert len(subsegments) == 1
    subsegment = subsegments[0]
    assert subsegment.name == 'dynamodb'
    assert len(subsegment.subsegments) == 0
    assert subsegment.error

    aws_meta = subsegment.aws
    assert aws_meta['region'] == 'us-west-2'
    assert aws_meta['operation'] == 'DescribeTable'
    assert aws_meta['table_name'] == 'mytable'


def test_empty_response():
    from aws_xray_sdk.ext.pynamodb.patch import pynamodb_meta_processor
    subsegment = xray_recorder.begin_subsegment('test')

    class TempReq:
        def __init__(self):
            self.headers = {'X-Amz-Target': 'ddb.ListTables'.encode('utf-8')}
            self.url = 'ddb.us-west-2'
            self.body = '{}'.encode('utf-8')

    prepared_request = TempReq()
    args = [prepared_request]

    pynamodb_meta_processor(wrapped=None, instance=None, args=args,
                            kwargs=None, return_value=None,
                            exception=None, subsegment=subsegment,
                            stack=None)

    aws_meta = subsegment.aws
    assert aws_meta['region'] == 'us-west-2'
    assert aws_meta['operation'] == 'ListTables'


def test_only_dynamodb_calls_are_traced():
    """Test only a single subsegment is created for other AWS services.

    As the pynamodb patch applies the botocore patch as well, we need
    to ensure that only one subsegment is created for all calls not
    made by PynamoDB. As PynamoDB calls botocore differently than the
    botocore patch expects we also just get a single subsegment per
    PynamoDB call.
    """
    session = botocore.session.get_session()
    s3 = session.create_client('s3', region_name='us-west-2',
                               config=Config(signature_version=UNSIGNED))
    try:
        s3.get_bucket_location(Bucket='mybucket')
    except ClientError:
        pass

    subsegments = xray_recorder.current_segment().subsegments
    assert len(subsegments) == 1
    assert subsegments[0].name == 's3'
    assert len(subsegments[0].subsegments) == 0