import boto3

from moto import mock_aws


@mock_aws
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]
    assert describe["AvailabilityZones"] == ["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"]
    )
    assert resp["AvailabilityZones"] == [
        "eu-north-1a",
        "eu-north-1b",
        "eu-north-1c",
        "eu-north-1d",
    ]

    describe = client.describe_load_balancers(LoadBalancerNames=["my-lb"])[
        "LoadBalancerDescriptions"
    ][0]
    assert describe["AvailabilityZones"] == [
        "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"]
    )
    assert resp["AvailabilityZones"] == ["eu-north-1a", "eu-north-1d"]

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