Add support for organizations and organizational units for EC2 image builder distribution configuration

internal/service/imagebuilder/distribution_configuration.go

@@ -89,6 +89,22 @@ func ResourceDistributionConfiguration() *schema.Resource {
 														ValidateFunc: verify.ValidAccountID,
 													},
 												},
+												"organization_arns": {
+													Type:     schema.TypeSet,
+													Optional: true,
+													Elem: &schema.Schema{
+														Type:         schema.TypeString,
+														ValidateFunc: verify.ValidARN,
+													},
+												},
+												"organizational_unit_arns": {
+													Type:     schema.TypeSet,
+													Optional: true,
+													Elem: &schema.Schema{
+														Type:         schema.TypeString,
+														ValidateFunc: verify.ValidARN,
+													},
+												},
 											},
 										},
 									},
@@ -387,6 +403,14 @@ func expandLaunchPermissionConfiguration(tfMap map[string]interface{}) *imagebui
 		apiObject.UserGroups = flex.ExpandStringSet(v)
 	}
 
+	if v, ok := tfMap["organization_arns"].(*schema.Set); ok && v.Len() > 0 {
+		apiObject.OrganizationArns = flex.ExpandStringSet(v)
+	}
+
+	if v, ok := tfMap["organizational_unit_arns"].(*schema.Set); ok && v.Len() > 0 {
+		apiObject.OrganizationalUnitArns = flex.ExpandStringSet(v)
+	}
+
 	return apiObject
 }
 
@@ -478,6 +502,13 @@ func flattenLaunchPermissionConfiguration(apiObject *imagebuilder.LaunchPermissi
 	if v := apiObject.UserIds; v != nil {
 		tfMap["user_ids"] = aws.StringValueSlice(v)
 	}
+	if v := apiObject.OrganizationArns; v != nil {
+		tfMap["organization_arns"] = aws.StringValueSlice(v)
+	}
+
+	if v := apiObject.OrganizationalUnitArns; v != nil {
+		tfMap["organizational_unit_arns"] = aws.StringValueSlice(v)
+	}
 
 	return tfMap
 }

internal/service/imagebuilder/distribution_configuration_data_source.go

@@ -85,6 +85,20 @@ func DataSourceDistributionConfiguration() *schema.Resource {
 											Type: schema.TypeString,
 										},
 									},
+									"organization_arns": {
+										Type:     schema.TypeSet,
+										Computed: true,
+										Elem: &schema.Schema{
+											Type: schema.TypeString,
+										},
+									},
+									"organizational_unit_arns": {
+										Type:     schema.TypeSet,
+										Computed: true,
+										Elem: &schema.Schema{
+											Type: schema.TypeString,
+										},
+									},
 								},
 							},
 						},

internal/service/imagebuilder/distribution_configuration_test.go

@@ -318,6 +318,69 @@ func TestAccImageBuilderDistributionConfiguration_DistributionAMIDistributionLau
 	})
 }
 
+func TestAccImageBuilderDistributionConfiguration_DistributionAMIDistributionLaunchPermission_organizationArns(t *testing.T) {
+	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
+	organizationResourceName := "aws_organizations_organization.test"
+	resourceName := "aws_imagebuilder_distribution_configuration.test"
+
+	resource.Test(t, resource.TestCase{
+		PreCheck: func() {
+			acctest.PreCheck(t)
+			acctest.PreCheckOrganizationsAccount(t)
+		},
+		ErrorCheck:        acctest.ErrorCheck(t, imagebuilder.EndpointsID),
+		ProviderFactories: acctest.ProviderFactories,
+		CheckDestroy:      testAccCheckDistributionConfigurationDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccDistributionConfigurationDistributionAMIDistributionConfigurationLaunchPermissionOrganizationArnsConfig(rName),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckDistributionConfigurationExists(resourceName),
+					resource.TestCheckResourceAttr(resourceName, "distribution.#", "1"),
+					resource.TestCheckTypeSetElemAttrPair(resourceName, "distribution.*.ami_distribution_configuration.0.launch_permission.0.organization_arns.*", organizationResourceName, "arn"),
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
+		},
+	})
+}
+
+func TestAccImageBuilderDistributionConfiguration_DistributionAMIDistributionLaunchPermission_organizationalUnitArns(t *testing.T) {
+	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
+	organizationalUnitResourceName := "aws_organizations_organizational_unit.test"
+
+	resourceName := "aws_imagebuilder_distribution_configuration.test"
+
+	resource.Test(t, resource.TestCase{
+		PreCheck: func() {
+			acctest.PreCheck(t)
+			acctest.PreCheckOrganizationsAccount(t)
+		},
+		ErrorCheck:        acctest.ErrorCheck(t, imagebuilder.EndpointsID),
+		ProviderFactories: acctest.ProviderFactories,
+		CheckDestroy:      testAccCheckDistributionConfigurationDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccDistributionConfigurationDistributionAMIDistributionConfigurationLaunchPermissionOrganizationalUnitArnsConfig(rName),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckDistributionConfigurationExists(resourceName),
+					resource.TestCheckResourceAttr(resourceName, "distribution.#", "1"),
+					resource.TestCheckTypeSetElemAttrPair(resourceName, "distribution.*.ami_distribution_configuration.0.launch_permission.0.organizational_unit_arns.*", organizationalUnitResourceName, "arn"),
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
+		},
+	})
+}
+
 func TestAccImageBuilderDistributionConfiguration_DistributionAMIDistribution_name(t *testing.T) {
 	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
 	resourceName := "aws_imagebuilder_distribution_configuration.test"
@@ -704,6 +767,53 @@ resource "aws_imagebuilder_distribution_configuration" "test" {
 `, rName, userId)
 }
 
+func testAccDistributionConfigurationDistributionAMIDistributionConfigurationLaunchPermissionOrganizationArnsConfig(rName string) string {
+	return fmt.Sprintf(`
+data "aws_region" "current" {}
+data "aws_partition" "current" {}
+
+resource "aws_organizations_organization" "test" {}
+
+resource "aws_imagebuilder_distribution_configuration" "test" {
+  name = %[1]q
+  distribution {
+    ami_distribution_configuration {
+      launch_permission {
+        organization_arns = [aws_organizations_organization.test.arn]
+      }
+    }
+    region = data.aws_region.current.name
+  }
+}
+`, rName)
+}
+
+func testAccDistributionConfigurationDistributionAMIDistributionConfigurationLaunchPermissionOrganizationalUnitArnsConfig(rName string) string {
+	return fmt.Sprintf(`
+data "aws_region" "current" {}
+data "aws_partition" "current" {}
+
+resource "aws_organizations_organization" "test" {}
+
+resource "aws_organizations_organizational_unit" "test" {
+	name = %[1]q
+	parent_id = aws_organizations_organization.test.roots[0].id
+}
+
+resource "aws_imagebuilder_distribution_configuration" "test" {
+	name = %[1]q
+	distribution {
+	  ami_distribution_configuration {
+		launch_permission {
+		  organizational_unit_arns = [aws_organizations_organizational_unit.test.arn]
+		}
+	  }
+	  region = data.aws_region.current.name
+	}
+  }
+  `, rName)
+}
+
 func testAccDistributionConfigurationDistributionAMIDistributionConfigurationNameConfig(rName string, name string) string {
 	return fmt.Sprintf(`
 data "aws_region" "current" {}