From 6982eed2b4474838eb74d77e2d80a04838ac5df9 Mon Sep 17 00:00:00 2001
From: Tong Liu <tongl@vmware.com>
Date: Fri, 24 Aug 2018 17:24:37 +0000
Subject: [PATCH] Expose allocation pool to router creation

Recently, platform exposes new API to create LR with allocation
profile so that user can specify allocation pool for LoadBalancer.

Also update the feature support util since this feature is only
available on version 2.3 and after. This feature can be enabled
based on nsx version.

e.g. Create LR with LB SMALL allocation pool
  allocation_pool = {
      "allocation_pool_type": "LoadBalancerAllocationPool",
      "allocation_size": "SMALL"}
  create(..., allocation_pool=allocation_pool)

Change-Id: I55cd17a81b3ff6705345b170d15f26fc9a5b0f74
---
 vmware_nsxlib/tests/unit/v3/test_resources.py | 12 ++++++++++--
 vmware_nsxlib/v3/__init__.py                  |  6 ++++++
 vmware_nsxlib/v3/core_resources.py            |  5 ++++-
 vmware_nsxlib/v3/nsx_constants.py             |  1 +
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/vmware_nsxlib/tests/unit/v3/test_resources.py b/vmware_nsxlib/tests/unit/v3/test_resources.py
index d57068ce..3b3ec69f 100644
--- a/vmware_nsxlib/tests/unit/v3/test_resources.py
+++ b/vmware_nsxlib/tests/unit/v3/test_resources.py
@@ -640,15 +640,23 @@ class LogicalRouterTestCase(BaseTestResource):
         tier0_router = True
         description = 'dummy'
         tz_id = 'tz_id'
+        allocation_pool = {
+            'allocation_pool_type': 'LoadBalancerAllocationPool',
+            'allocation_size': 'SMALL'
+        }
         router.create(fake_router['display_name'], None, None, tier0_router,
-                      description=description, transport_zone_id=tz_id)
+                      description=description, transport_zone_id=tz_id,
+                      allocation_pool=allocation_pool)
 
         data = {
             'display_name': fake_router['display_name'],
             'router_type': 'TIER0' if tier0_router else 'TIER1',
             'tags': None,
             'description': description,
-            'advanced_config': {'transport_zone_id': tz_id}
+            'advanced_config': {'transport_zone_id': tz_id},
+            'allocation_profile': {
+                'allocation_pool': allocation_pool
+            }
         }
 
         test_client.assert_json_call(
diff --git a/vmware_nsxlib/v3/__init__.py b/vmware_nsxlib/v3/__init__.py
index b98312d7..bd981e78 100644
--- a/vmware_nsxlib/v3/__init__.py
+++ b/vmware_nsxlib/v3/__init__.py
@@ -328,6 +328,12 @@ class NsxLib(NsxLibBase):
             if (feature == nsx_constants.FEATURE_ICMP_STRICT):
                 return True
 
+        if (version.LooseVersion(self.get_version()) >=
+                version.LooseVersion(nsx_constants.NSX_VERSION_2_3_0)):
+            # Features available since 2.3
+            if (feature == nsx_constants.FEATURE_ALLOCATION_POOL):
+                return True
+
         if (version.LooseVersion(self.get_version()) >=
             version.LooseVersion(nsx_constants.NSX_VERSION_2_2_0)):
             # Features available since 2.2
diff --git a/vmware_nsxlib/v3/core_resources.py b/vmware_nsxlib/v3/core_resources.py
index 7b42d626..2b3c762f 100644
--- a/vmware_nsxlib/v3/core_resources.py
+++ b/vmware_nsxlib/v3/core_resources.py
@@ -650,7 +650,7 @@ class NsxLibLogicalRouter(utils.NsxLibApiBase):
         return self.client.get(resource)
 
     def create(self, display_name, tags, edge_cluster_uuid=None, tier_0=False,
-               description=None, transport_zone_id=None):
+               description=None, transport_zone_id=None, allocation_pool=None):
         # TODO(salv-orlando): If possible do not manage edge clusters
         # in the main plugin logic.
         router_type = (nsx_constants.ROUTER_TYPE_TIER0 if tier_0 else
@@ -665,6 +665,9 @@ class NsxLibLogicalRouter(utils.NsxLibApiBase):
         if transport_zone_id:
             body['advanced_config'] = {
                 'transport_zone_id': transport_zone_id}
+        if allocation_pool:
+            body['allocation_profile'] = {
+                'allocation_pool': allocation_pool}
         return self.client.create(self.get_path(), body=body)
 
     def delete(self, lrouter_id, force=False):
diff --git a/vmware_nsxlib/v3/nsx_constants.py b/vmware_nsxlib/v3/nsx_constants.py
index eaa40e2f..467a3db7 100644
--- a/vmware_nsxlib/v3/nsx_constants.py
+++ b/vmware_nsxlib/v3/nsx_constants.py
@@ -146,3 +146,4 @@ FEATURE_ROUTER_TRANSPORT_ZONE = 'Router Transport Zone'
 FEATURE_NO_DNAT_NO_SNAT = 'No DNAT/No SNAT'
 FEATURE_ENS_WITH_SEC = 'ENS with security'
 FEATURE_ICMP_STRICT = 'Strict list of supported ICMP types and codes'
+FEATURE_ALLOCATION_POOL = 'Router Allocation Pool'