import boto3

from moto import mock_aws


@mock_aws
def test_elb_attach_load_balancer_to_subnets():
    ec2 = boto3.resource("ec2", region_name="us-east-1")
    vpc = ec2.create_vpc(CidrBlock="172.28.7.0/24", InstanceTenancy="default")
    subnet1 = ec2.create_subnet(VpcId=vpc.id, CidrBlock="172.28.7.192/26")
    subnet2 = ec2.create_subnet(VpcId=vpc.id, CidrBlock="172.28.7.191/26")
    client = boto3.client("elb", region_name="us-east-1")
    client.create_load_balancer(
        LoadBalancerName="my-lb",
        Listeners=[{"Protocol": "tcp", "LoadBalancerPort": 80, "InstancePort": 8080}],
        Subnets=[subnet1.id],
    )

    client.attach_load_balancer_to_subnets(
        LoadBalancerName="my-lb", Subnets=[subnet2.id]
    )

    lb = client.describe_load_balancers()["LoadBalancerDescriptions"][0]

    assert len(lb["Subnets"]) == 2
    assert subnet1.id in lb["Subnets"]
    assert subnet2.id in lb["Subnets"]


@mock_aws
def test_elb_detach_load_balancer_to_subnets():
    ec2 = boto3.resource("ec2", region_name="us-east-1")
    vpc = ec2.create_vpc(CidrBlock="172.28.7.0/24", InstanceTenancy="default")
    subnet1 = ec2.create_subnet(VpcId=vpc.id, CidrBlock="172.28.7.192/26")
    subnet2 = ec2.create_subnet(VpcId=vpc.id, CidrBlock="172.28.7.191/26")
    client = boto3.client("elb", region_name="us-east-1")
    client.create_load_balancer(
        LoadBalancerName="my-lb",
        Listeners=[{"Protocol": "tcp", "LoadBalancerPort": 80, "InstancePort": 8080}],
        Subnets=[subnet1.id, subnet2.id],
    )

    client.detach_load_balancer_from_subnets(
        LoadBalancerName="my-lb", Subnets=[subnet1.id]
    )

    lb = client.describe_load_balancers()["LoadBalancerDescriptions"][0]

    assert len(lb["Subnets"]) == 1
    assert subnet2.id in lb["Subnets"]
