commit 0f90a705966dc2bfb71c0fbed83eb64b846eb2d6
Author: Erica Liu <ericaliu@vmware.com>
Date:   Tue Oct 13 17:07:36 2020 -0700

    Adding support for customzing ep value in segement creation
    
    In case of none default enforcement point value, the Segment
    create might fail, because it is hard coded to use default ep
    for querying transport zone, if specified. This change add
    support for creating segment with transport zone in none default
    enforcement point
    
    Change-Id: Id122f9591c2bded5edc43fad514e6e1e9e6a9fa3
    (cherry picked from commit ba0994042982c81dede2f7e7800bc479ca23d962)

diff --git a/vmware_nsxlib/tests/unit/v3/policy/test_resources.py b/vmware_nsxlib/tests/unit/v3/policy/test_resources.py
index 4cc19af..a34e122 100644
--- a/vmware_nsxlib/tests/unit/v3/policy/test_resources.py
+++ b/vmware_nsxlib/tests/unit/v3/policy/test_resources.py
@@ -4093,7 +4093,8 @@ class TestPolicySegment(NsxPolicyLibTestCase):
 
     def _test_create(self, tier1_id=None, tier0_id=None, mdproxy=None,
                      dhcp_server=None, admin_state=None,
-                     ip_pool_id='external-ip-pool', ls_id=None):
+                     ip_pool_id='external-ip-pool', ls_id=None,
+                     tz_id=None, ep_id=None):
         name = 'test'
         description = 'desc'
         subnets = [core_defs.Subnet(gateway_address="2.2.2.0/24")]
@@ -4118,6 +4119,11 @@ class TestPolicySegment(NsxPolicyLibTestCase):
         if ls_id:
             kwargs['ls_id'] = ls_id
 
+        if tz_id:
+            kwargs['transport_zone_id'] = tz_id
+        if ep_id:
+            kwargs['ep_id'] = ep_id
+
         with mock.patch.object(self.policy_api,
                                "create_or_update") as api_call:
             result = self.resourceApi.create_or_overwrite(name, **kwargs)
@@ -4171,6 +4177,12 @@ class TestPolicySegment(NsxPolicyLibTestCase):
     def test_create_with_ls_id(self):
         self._test_create(ls_id='lsid1')
 
+    def test_create_with_transport_zone_id(self):
+        self._test_create(tz_id='tz_id1', ep_id='ep_id1')
+
+    def test_create_with_transport_zone_id_and_default_ep(self):
+        self._test_create(tz_id='tz_id1')
+
     def test_delete(self):
         segment_id = '111'
         with mock.patch.object(self.policy_api, "delete") as api_call:
diff --git a/vmware_nsxlib/v3/policy/core_defs.py b/vmware_nsxlib/v3/policy/core_defs.py
index 2d7c2e4..45426b2 100644
--- a/vmware_nsxlib/v3/policy/core_defs.py
+++ b/vmware_nsxlib/v3/policy/core_defs.py
@@ -996,7 +996,8 @@ class SegmentDef(BaseSegmentDef):
             if self.get_attr('transport_zone_id'):
                 tz = TransportZoneDef(
                     tz_id=self.get_attr('transport_zone_id'),
-                    ep_id=constants.DEFAULT_ENFORCEMENT_POINT,
+                    ep_id=self.get_attr(
+                        'ep_id') or constants.DEFAULT_ENFORCEMENT_POINT,
                     tenant=self.get_tenant())
                 path = tz.get_resource_full_path()
             self._set_attr_if_specified(body, 'transport_zone_id',
diff --git a/vmware_nsxlib/v3/policy/core_resources.py b/vmware_nsxlib/v3/policy/core_resources.py
index a7f8e63..05ed952 100644
--- a/vmware_nsxlib/v3/policy/core_resources.py
+++ b/vmware_nsxlib/v3/policy/core_resources.py
@@ -1947,6 +1947,7 @@ class NsxPolicySegmentApi(NsxPolicyResourceBase):
                             dhcp_server_config_id=IGNORE,
                             admin_state=IGNORE,
                             ls_id=IGNORE,
+                            ep_id=IGNORE,
                             tags=IGNORE,
                             tenant=constants.POLICY_INFRA_TENANT):
 
@@ -1971,6 +1972,7 @@ class NsxPolicySegmentApi(NsxPolicyResourceBase):
             dhcp_server_config_id=dhcp_server_config_id,
             admin_state=admin_state,
             ls_id=ls_id,
+            ep_id=ep_id,
             tags=tags,
             tenant=tenant)
         self._create_or_store(segment_def)