Consolidate 'findIPAMPoolAllocationsForEIP' and 'findIPAMPoolAllocationsForVPC' into 'findIPAMPoolAllocationForResource'.

Author: ewbankkit

internal/service/ec2/ec2_eip.go

@@ -16,12 +16,10 @@ import (
 	awstypes "github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/hashicorp/aws-sdk-go-base/v2/tfawserr"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
-	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/retry"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-provider-aws/internal/conns"
 	"github.com/hashicorp/terraform-provider-aws/internal/enum"
 	"github.com/hashicorp/terraform-provider-aws/internal/errs/sdkdiag"
-	tfslices "github.com/hashicorp/terraform-provider-aws/internal/slices"
 	tftags "github.com/hashicorp/terraform-provider-aws/internal/tags"
 	"github.com/hashicorp/terraform-provider-aws/internal/tfresource"
 	"github.com/hashicorp/terraform-provider-aws/names"
@@ -339,7 +337,7 @@ func resourceEIPDelete(ctx context.Context, d *schema.ResourceData, meta any) di
 			timeout = 10 * time.Minute // IPAM eventual consistency
 		)
 		_, err := tfresource.RetryUntilNotFound(ctx, timeout, func(ctx context.Context) (any, error) {
-			return findIPAMPoolAllocationsForEIP(ctx, conn, ipamPoolID, d.Get("allocation_id").(string))
+			return findIPAMPoolAllocationForResource(ctx, conn, ipamPoolID, d.Get("allocation_id").(string))
 		})
 
 		if err != nil {
@@ -430,25 +428,3 @@ func disassociateEIP(ctx context.Context, conn *ec2.Client, associationID string
 func eipARN(ctx context.Context, c *conns.AWSClient, allocationID string) string {
 	return c.RegionalARN(ctx, names.EC2, "elastic-ip/"+allocationID)
 }
-
-func findIPAMPoolAllocationsForEIP(ctx context.Context, conn *ec2.Client, ipamPoolID, eipAllocationID string) ([]awstypes.IpamPoolAllocation, error) {
-	input := ec2.GetIpamPoolAllocationsInput{
-		IpamPoolId: aws.String(ipamPoolID),
-	}
-
-	output, err := findIPAMPoolAllocations(ctx, conn, &input)
-
-	if err != nil {
-		return nil, err
-	}
-
-	output = tfslices.Filter(output, func(v awstypes.IpamPoolAllocation) bool {
-		return v.ResourceType == awstypes.IpamPoolAllocationResourceTypeEip && aws.ToString(v.ResourceId) == eipAllocationID
-	})
-
-	if len(output) == 0 {
-		return nil, &retry.NotFoundError{}
-	}
-
-	return output, nil
-}

internal/service/ec2/exports_test.go

@@ -162,7 +162,7 @@ var (
 	FindHostByID                                                = findHostByID
 	FindIPAMByID                                                = findIPAMByID
 	FindIPAMPoolAllocationByTwoPartKey                          = findIPAMPoolAllocationByTwoPartKey
-	FindIPAMPoolAllocationsForVPC                               = findIPAMPoolAllocationsForVPC
+	FindIPAMPoolAllocationForResource                           = findIPAMPoolAllocationForResource
 	FindIPAMPoolByID                                            = findIPAMPoolByID
 	FindIPAMPoolCIDRByTwoPartKey                                = findIPAMPoolCIDRByTwoPartKey
 	FindIPAMResourceDiscoveryAssociationByID                    = findIPAMResourceDiscoveryAssociationByID

internal/service/ec2/find.go

@@ -4188,26 +4188,14 @@ func findIPAMPoolAllocationsByIPAMPoolIDAndResourceID(ctx context.Context, conn
 	}), nil
 }
 
-func findIPAMPoolAllocationsForVPC(ctx context.Context, conn *ec2.Client, poolID, vpcID string) ([]awstypes.IpamPoolAllocation, error) {
-	input := ec2.GetIpamPoolAllocationsInput{
-		IpamPoolId: aws.String(poolID),
-	}
-
-	output, err := findIPAMPoolAllocations(ctx, conn, &input)
+func findIPAMPoolAllocationForResource(ctx context.Context, conn *ec2.Client, ipamPoolID, resourceID string) (*awstypes.IpamPoolAllocation, error) {
+	output, err := findIPAMPoolAllocationsByIPAMPoolIDAndResourceID(ctx, conn, ipamPoolID, resourceID)
 
 	if err != nil {
 		return nil, err
 	}
 
-	output = tfslices.Filter(output, func(v awstypes.IpamPoolAllocation) bool {
-		return v.ResourceType == awstypes.IpamPoolAllocationResourceTypeVpc && aws.ToString(v.ResourceId) == vpcID
-	})
-
-	if len(output) == 0 {
-		return nil, &retry.NotFoundError{}
-	}
-
-	return output, nil
+	return tfresource.AssertFirstValueResult(output)
 }
 
 func findIPAMPoolCIDR(ctx context.Context, conn *ec2.Client, input *ec2.GetIpamPoolCidrsInput) (*awstypes.IpamPoolCidr, error) {

internal/service/ec2/vpc_.go

@@ -483,7 +483,7 @@ func resourceVPCDelete(ctx context.Context, d *schema.ResourceData, meta any) di
 			timeout = 35 * time.Minute // IPAM eventual consistency. It can take ~30 min to release allocations.
 		)
 		_, err := tfresource.RetryUntilNotFound(ctx, timeout, func(ctx context.Context) (any, error) {
-			return findIPAMPoolAllocationsForVPC(ctx, conn, ipamPoolID, d.Id())
+			return findIPAMPoolAllocationForResource(ctx, conn, ipamPoolID, d.Id())
 		})
 
 		if err != nil {

internal/service/ec2/vpc_ipv4_cidr_block_association_test.go

@@ -277,7 +277,7 @@ func testAccCheckVPCIPv4CIDRBlockAssociationWaitVPCIPAMPoolAllocationDeleted(ctx
 			timeout = 35 * time.Minute // IPAM eventual consistency. It can take ~30 min to release allocations.
 		)
 		_, err := tfresource.RetryUntilNotFound(ctx, timeout, func(ctx context.Context) (any, error) {
-			return tfec2.FindIPAMPoolAllocationsForVPC(ctx, conn, rsIPAMPool.Primary.ID, rsVPC.Primary.ID)
+			return tfec2.FindIPAMPoolAllocationForResource(ctx, conn, rsIPAMPool.Primary.ID, rsVPC.Primary.ID)
 		})
 
 		return err