From 562bde334a85086f01f0812fdefd6f51b45abda3 Mon Sep 17 00:00:00 2001
From: Edward Hope-Morley <edward.hope-morley@canonical.com>
Date: Thu, 19 Nov 2015 17:58:38 +0000
Subject: [PATCH] add pool create pg_num override support

---
 hooks/ceph_broker.py           | 11 ++++++++++-
 unit_tests/test_ceph_broker.py | 21 +++++++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/hooks/ceph_broker.py b/hooks/ceph_broker.py
index e162dcb..e9d5e5a 100644
--- a/hooks/ceph_broker.py
+++ b/hooks/ceph_broker.py
@@ -83,12 +83,21 @@ def process_requests_v1(reqs):
                 log(msg, level=ERROR)
                 return {'exit-code': 1, 'stderr': msg}
 
+            # Mandatory params
             pool = params['pool']
             replicas = params['replicas']
+
+            # Optional params
+            pg_num = req.get('pg_num')
+            if pg_num:
+                # Ensure string
+                pg_num = str(pg_num)
+
             if not pool_exists(service=svc, name=pool):
                 log("Creating pool '%s' (replicas=%s)" % (pool, replicas),
                     level=INFO)
-                create_pool(service=svc, name=pool, replicas=replicas)
+                create_pool(service=svc, name=pool, replicas=replicas,
+                            pg_num=pg_num)
             else:
                 log("Pool '%s' already exists - skipping create" % (pool),
                     level=DEBUG)
diff --git a/unit_tests/test_ceph_broker.py b/unit_tests/test_ceph_broker.py
index fc69817..93ca98e 100644
--- a/unit_tests/test_ceph_broker.py
+++ b/unit_tests/test_ceph_broker.py
@@ -53,7 +53,24 @@ class CephBrokerTestCase(unittest.TestCase):
         rc = ceph_broker.process_requests(reqs)
         mock_pool_exists.assert_called_with(service='admin', name='foo')
         mock_create_pool.assert_called_with(service='admin', name='foo',
-                                            replicas=3)
+                                            replicas=3, pg_num=None)
+        self.assertEqual(json.loads(rc), {'exit-code': 0})
+
+    @mock.patch('ceph_broker.create_pool')
+    @mock.patch('ceph_broker.pool_exists')
+    @mock.patch('ceph_broker.log')
+    def test_process_requests_create_pool_w_pg_num(self, mock_log,
+                                                   mock_pool_exists,
+                                                   mock_create_pool):
+        mock_pool_exists.return_value = False
+        reqs = json.dumps({'api-version': 1,
+                           'ops': [{'op': 'create-pool', 'name':
+                                    'foo', 'replicas': 3,
+                                    'pg_num': 100}]})
+        rc = ceph_broker.process_requests(reqs)
+        mock_pool_exists.assert_called_with(service='admin', name='foo')
+        mock_create_pool.assert_called_with(service='admin', name='foo',
+                                            replicas=3, pg_num='100')
         self.assertEqual(json.loads(rc), {'exit-code': 0})
 
     @mock.patch('ceph_broker.create_pool')
@@ -84,7 +101,7 @@ class CephBrokerTestCase(unittest.TestCase):
         rc = ceph_broker.process_requests(reqs)
         mock_pool_exists.assert_called_with(service='admin', name='foo')
         mock_create_pool.assert_called_with(service='admin', name='foo',
-                                            replicas=3)
+                                            replicas=3, pg_num=None)
         self.assertEqual(json.loads(rc)['exit-code'], 0)
         self.assertEqual(json.loads(rc)['request-id'], '1ef5aede')