import boto3

from moto import mock_aws


@mock_aws
def test_set_subnets_errors():
    client = boto3.client("elbv2", region_name="us-east-1")
    ec2 = boto3.resource("ec2", region_name="us-east-1")

    security_group = ec2.create_security_group(
        GroupName="a-security-group", Description="First One"
    )
    vpc = ec2.create_vpc(CidrBlock="172.28.7.0/24", InstanceTenancy="default")
    subnet1 = ec2.create_subnet(
        VpcId=vpc.id, CidrBlock="172.28.7.0/26", AvailabilityZone="us-east-1a"
    )
    subnet2 = ec2.create_subnet(
        VpcId=vpc.id, CidrBlock="172.28.7.64/26", AvailabilityZone="us-east-1b"
    )
    subnet3 = ec2.create_subnet(
        VpcId=vpc.id, CidrBlock="172.28.7.192/26", AvailabilityZone="us-east-1c"
    )

    response = client.create_load_balancer(
        Name="my-lb",
        Subnets=[subnet1.id, subnet2.id],
        SecurityGroups=[security_group.id],
        Scheme="internal",
        Tags=[{"Key": "key_name", "Value": "a_value"}],
    )
    arn = response["LoadBalancers"][0]["LoadBalancerArn"]

    client.set_subnets(
        LoadBalancerArn=arn, Subnets=[subnet1.id, subnet2.id, subnet3.id]
    )

    resp = client.describe_load_balancers(LoadBalancerArns=[arn])
    assert len(resp["LoadBalancers"][0]["AvailabilityZones"]) == 3


@mock_aws
def test_set_subnets__mapping():
    client = boto3.client("elbv2", region_name="us-east-1")
    ec2 = boto3.resource("ec2", region_name="us-east-1")

    security_group = ec2.create_security_group(
        GroupName="a-security-group", Description="First One"
    )
    vpc = ec2.create_vpc(CidrBlock="172.28.7.0/24", InstanceTenancy="default")
    subnet1 = ec2.create_subnet(
        VpcId=vpc.id, CidrBlock="172.28.7.0/26", AvailabilityZone="us-east-1a"
    )
    subnet2 = ec2.create_subnet(
        VpcId=vpc.id, CidrBlock="172.28.7.64/26", AvailabilityZone="us-east-1b"
    )

    response = client.create_load_balancer(
        Name="my-lb",
        Subnets=[subnet1.id, subnet2.id],
        SecurityGroups=[security_group.id],
        Scheme="internal",
        Tags=[{"Key": "key_name", "Value": "a_value"}],
    )
    arn = response["LoadBalancers"][0]["LoadBalancerArn"]

    client.set_subnets(
        LoadBalancerArn=arn,
        SubnetMappings=[{"SubnetId": subnet1.id}, {"SubnetId": subnet2.id}],
    )

    resp = client.describe_load_balancers(LoadBalancerArns=[arn])
    a_zones = resp["LoadBalancers"][0]["AvailabilityZones"]
    assert len(a_zones) == 2
    assert {"ZoneName": "us-east-1a", "SubnetId": subnet1.id} in a_zones
    assert {"ZoneName": "us-east-1b", "SubnetId": subnet2.id} in a_zones