import pytest
from unittest.mock import patch

import boto3
import moto
from psutil._common import sdiskpart
import unittest
from zenyatta.aws.ec2 import EBSVolume, next_letter, get_free_device_from_self, create_ec2_instance
from tests.util import get_ec2_instance, get_task_instance, get_dag


class EBSTestCase(unittest.TestCase):

    @moto.mock_ec2
    @moto.mock_sts
    def test_ebs_create(self):
        with pytest.raises(AssertionError):
            EBSVolume(size=500)

        with pytest.raises(AssertionError):
            EBSVolume(size=500, volume_type='gp2')
        volume = EBSVolume(size=500, volume_type='gp2', availability_zone='us-east-1',
                           role_arn="arn:aws:iam::111111111111:role/airflow-dev")
        response = volume.create()
        assert response is not None
        assert response['ResponseMetadata']['HTTPStatusCode'] == 200
        assert response['State'] == 'creating'

    @moto.mock_ec2
    def test_ebs_destroy(self):
        pass

    @moto.mock_ec2
    @moto.mock_sts
    def test_ebs_attach_detach(self):
        # need an instance first
        ec2 = boto3.resource('ec2', region_name='us-west-2')
        instance = ec2.create_instances(ImageId='ami-1234', MinCount=1, MaxCount=1).pop()
        ebs = EBSVolume(size=500, volume_type='gp2', availability_zone='us-west-2a',
                        role_arn="arn:aws:iam::111111111111:role/airflow-dev")
        ebs.create()
        device = '/dev/xvdf'
        response = ebs.attach(instance.instance_id, device)
        assert response['ResponseMetadata']['HTTPStatusCode'] == 200
        assert response['State'] == 'attaching'
        attach_status = ebs.attach_status(instance.instance_id)
        assert attach_status == 'attached'
        # now detach
        response = ebs.detach(instance.instance_id, device)
        assert response['ResponseMetadata']['HTTPStatusCode'] == 200
        assert response['State'] == 'detaching'
        attach_status = ebs.attach_status(instance.instance_id)
        assert attach_status == 'detached'

    @moto.mock_ec2
    @moto.mock_sts
    def test_ebs_available(self):
        ebs = EBSVolume(size=500, volume_type='gp2', availability_zone='us-west-2a',
                        role_arn="arn:aws:iam::111111111111:role/airflow-dev")
        ebs.create()
        assert ebs.available() is True


@patch('zenyatta.aws.ec2.psutil.disk_partitions',
       return_value=[
           sdiskpart(device='/dev/xvda1', mountpoint='/', fstype='ext4',
                     opts='rw,relatime,discard,data=ordered'),
           sdiskpart(device='/dev/xvdb', mountpoint='/', fstype='ext4',
                     opts='rw,relatime,discard,data=ordered'),
           sdiskpart(device='/dev/xvdc', mountpoint='/', fstype='ext4',
                     opts='rw,relatime,discard,data=ordered'),
       ])
def test_get_free_device_from_self(psutil_mock):
    device = get_free_device_from_self([])
    assert device == '/dev/xvdd'


@patch('zenyatta.aws.ec2.conf.get', return_value='arn:aws:iam::641044725657:instance-profile/zenyatta')
@patch('zenyatta.aws.conf.get', return_value='us-west-2')
@moto.mock_sts
@moto.mock_ec2
def test_create_ec2_instance(ec2_conf_mock, aws_conf_mock):
    instance = get_ec2_instance()
    ec2 = boto3.resource('ec2', region_name='us-west-2')
    vpc = ec2.create_vpc(CidrBlock="10.0.0.0/24")
    sg = vpc.create_security_group(GroupName="test-sg", Description="test")
    subnet = vpc.create_subnet(CidrBlock="10.0.0.0/25")
    new_instances = create_ec2_instance(image_id='ami-1234',
                                        instance_type=instance.instance_type,
                                        security_groups=[{'GroupId': sg.id}],
                                        min_count=1, max_count=1,
                                        task_instance=get_task_instance(),
                                        ts_nodash='20170105',
                                        subnet_id=subnet.id,
                                        availability_zone='us-west-2',
                                        role_arn='arn:aws:iam::641044725657:role/airflow')
    assert len(new_instances) == 1
    new_instance = ec2.Instance(new_instances.pop())
    assert new_instance.security_groups == [{'GroupId': sg.id, 'GroupName': 'test-sg'}]
    assert new_instance.instance_type == instance.instance_type
    assert new_instance.image_id == 'ami-1234'
    assert new_instance.state['Name'] == 'running'
    assert new_instance.state['Code'] == 16


def test_next_letter():
    assert 'c' == next_letter('a')
    assert 'c' == next_letter('z')
    assert 'z' == next_letter('y')
    with pytest.raises(AssertionError):
        next_letter('A')
