import boto3
import sure  # noqa # pylint: disable=unused-import

from moto import mock_elb, mock_ec2


@mock_elb
@mock_ec2
def test_enable_and_disable_availability_zones():
    client = boto3.client("elb", region_name="eu-north-1")
    ec2 = boto3.resource("ec2", region_name="eu-north-1")

    security_group = ec2.create_security_group(GroupName="sg01", Description="desc")

    client.create_load_balancer(
        LoadBalancerName="my-lb",
        Listeners=[
            {"Protocol": "tcp", "LoadBalancerPort": 80, "InstancePort": 8080},
            {"Protocol": "http", "LoadBalancerPort": 81, "InstancePort": 9000},
        ],
        AvailabilityZones=["eu-north-1a", "eu-north-1b"],
        Scheme="internal",
        SecurityGroups=[security_group.id],
    )

    describe = client.describe_load_balancers(LoadBalancerNames=["my-lb"])[
        "LoadBalancerDescriptions"
    ][0]
    describe["AvailabilityZones"].should.equal(["eu-north-1a", "eu-north-1b"])

    # enable more az's
    resp = client.enable_availability_zones_for_load_balancer(
        LoadBalancerName="my-lb", AvailabilityZones=["eu-north-1c", "eu-north-1d"]
    )
    resp.should.have.key("AvailabilityZones").equals(
        ["eu-north-1a", "eu-north-1b", "eu-north-1c", "eu-north-1d"]
    )

    describe = client.describe_load_balancers(LoadBalancerNames=["my-lb"])[
        "LoadBalancerDescriptions"
    ][0]
    describe["AvailabilityZones"].should.equal(
        ["eu-north-1a", "eu-north-1b", "eu-north-1c", "eu-north-1d"]
    )

    # disable some az's
    resp = client.disable_availability_zones_for_load_balancer(
        LoadBalancerName="my-lb", AvailabilityZones=["eu-north-1b", "eu-north-1c"]
    )
    resp.should.have.key("AvailabilityZones").equals(["eu-north-1a", "eu-north-1d"])

    describe = client.describe_load_balancers(LoadBalancerNames=["my-lb"])[
        "LoadBalancerDescriptions"
    ][0]
    describe["AvailabilityZones"].should.equal(["eu-north-1a", "eu-north-1d"])