From 01039079ff2ef28039e034b53563408ae061a055 Mon Sep 17 00:00:00 2001
From: Anna Khmelnitsky <akhmelnitsky@vmware.com>
Date: Mon, 14 Sep 2020 14:31:13 -0700
Subject: [PATCH] Expose vlan_ids attribute for overlay segments

This is to allow use case where VM tags packets and/or EVPN
---
 nsxt/resource_nsxt_policy_segment.go        | 5 +----
 nsxt/resource_nsxt_policy_segment_test.go   | 4 ++++
 nsxt/resource_nsxt_policy_vlan_segment.go   | 2 +-
 nsxt/segment_common.go                      | 5 +++--
 website/docs/r/policy_segment.html.markdown | 1 +
 5 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/nsxt/resource_nsxt_policy_segment.go b/nsxt/resource_nsxt_policy_segment.go
index 94ebde0dd..1d5dfa4f5 100644
--- a/nsxt/resource_nsxt_policy_segment.go
+++ b/nsxt/resource_nsxt_policy_segment.go
@@ -8,9 +8,6 @@ import (
 )
 
 func resourceNsxtPolicySegment() *schema.Resource {
-	segSchema := getPolicyCommonSegmentSchema()
-	delete(segSchema, "vlan_ids")
-
 	return &schema.Resource{
 		Create: resourceNsxtPolicySegmentCreate,
 		Read:   resourceNsxtPolicySegmentRead,
@@ -20,7 +17,7 @@ func resourceNsxtPolicySegment() *schema.Resource {
 			State: schema.ImportStatePassthrough,
 		},
 
-		Schema: segSchema,
+		Schema: getPolicyCommonSegmentSchema(false),
 	}
 }
 
diff --git a/nsxt/resource_nsxt_policy_segment_test.go b/nsxt/resource_nsxt_policy_segment_test.go
index 9f4a7674f..f25e1dfba 100644
--- a/nsxt/resource_nsxt_policy_segment_test.go
+++ b/nsxt/resource_nsxt_policy_segment_test.go
@@ -146,6 +146,7 @@ func TestAccResourceNsxtPolicySegment_updateAdvConfig(t *testing.T) {
 					resource.TestCheckResourceAttr(testResourceName, "subnet.0.cidr", "12.12.2.1/24"),
 					resource.TestCheckResourceAttr(testResourceName, "domain_name", "tftest.org"),
 					resource.TestCheckResourceAttr(testResourceName, "overlay_id", "1011"),
+					resource.TestCheckResourceAttr(testResourceName, "vlan_ids.#", "2"),
 					resource.TestCheckResourceAttr(testResourceName, "tag.#", "1"),
 					resource.TestCheckResourceAttr(testResourceName, "advanced_config.0.connectivity", "OFF"),
 					resource.TestCheckResourceAttr(testResourceName, "advanced_config.0.local_egress", "true"),
@@ -161,6 +162,7 @@ func TestAccResourceNsxtPolicySegment_updateAdvConfig(t *testing.T) {
 					resource.TestCheckResourceAttr(testResourceName, "subnet.0.cidr", "12.12.2.1/24"),
 					resource.TestCheckResourceAttr(testResourceName, "domain_name", "tftest.org"),
 					resource.TestCheckResourceAttr(testResourceName, "overlay_id", "1011"),
+					resource.TestCheckResourceAttr(testResourceName, "vlan_ids.#", "1"),
 					resource.TestCheckResourceAttr(testResourceName, "tag.#", "1"),
 					resource.TestCheckResourceAttr(testResourceName, "advanced_config.0.connectivity", "ON"),
 					resource.TestCheckResourceAttr(testResourceName, "advanced_config.0.local_egress", "false"),
@@ -524,6 +526,7 @@ resource "nsxt_policy_segment" "test" {
   description  = "Acceptance Test"
   domain_name  = "tftest.org"
   overlay_id   = 1011
+  vlan_ids     = ["101", "102"]
 
   transport_zone_path = data.nsxt_policy_transport_zone.test.path
 
@@ -552,6 +555,7 @@ resource "nsxt_policy_segment" "test" {
   description  = "Acceptance Test"
   domain_name  = "tftest.org"
   overlay_id   = 1011
+  vlan_ids     = ["101-104"]
 
   transport_zone_path = data.nsxt_policy_transport_zone.test.path
 
diff --git a/nsxt/resource_nsxt_policy_vlan_segment.go b/nsxt/resource_nsxt_policy_vlan_segment.go
index 5e5800626..144b78742 100644
--- a/nsxt/resource_nsxt_policy_vlan_segment.go
+++ b/nsxt/resource_nsxt_policy_vlan_segment.go
@@ -8,7 +8,7 @@ import (
 )
 
 func resourceNsxtPolicyVlanSegment() *schema.Resource {
-	segSchema := getPolicyCommonSegmentSchema()
+	segSchema := getPolicyCommonSegmentSchema(true)
 	delete(segSchema, "overlay_id")
 	delete(segSchema, "connectivity_path")
 
diff --git a/nsxt/segment_common.go b/nsxt/segment_common.go
index 640ca2d03..cef16493b 100644
--- a/nsxt/segment_common.go
+++ b/nsxt/segment_common.go
@@ -238,7 +238,7 @@ func getPolicySegmentSecurityProfilesSchema() *schema.Resource {
 	}
 }
 
-func getPolicyCommonSegmentSchema() map[string]*schema.Schema {
+func getPolicyCommonSegmentSchema(vlanRequired bool) map[string]*schema.Schema {
 	return map[string]*schema.Schema{
 		"nsx_id":       getNsxIDSchema(),
 		"path":         getPathSchema(),
@@ -297,7 +297,8 @@ func getPolicyCommonSegmentSchema() map[string]*schema.Schema {
 				Type:         schema.TypeString,
 				ValidateFunc: validateVLANIdOrRange,
 			},
-			Required: true,
+			Required: vlanRequired,
+			Optional: !vlanRequired,
 		},
 		"discovery_profile": {
 			Type:        schema.TypeList,
diff --git a/website/docs/r/policy_segment.html.markdown b/website/docs/r/policy_segment.html.markdown
index 5f6d9b340..cce949168 100644
--- a/website/docs/r/policy_segment.html.markdown
+++ b/website/docs/r/policy_segment.html.markdown
@@ -59,6 +59,7 @@ The following arguments are supported:
 * `connectivity_path` - (Optional) Policy path to the connecting Tier-0 or Tier-1.
 * `domain_name`- (Optional) DNS domain names.
 * `overlay_id` - (Optional) Overlay connectivity ID for this Segment.
+* `vlan_ids` - (Optional) List of VLAN IDs or ranges. Specifying vlan ids can be useful for overlay segments, f.e. for EVPN.
 * `transport_zone_path` - (Required) Policy path to the Overlay transport zone. This property is required if more than one overlay transport zone is defined, and none is marked as default.
 * `dhcp_config_path` - (Optional) Policy path to DHCP server or relay configuration to use for subnets configured on this segment. This attribute is supported with NSX 3.0.0 onwards.
 * `subnet` - (Optional) Subnet configuration block.