From c8df2602fd8f4f2cb7716e6283f3779c6895a479 Mon Sep 17 00:00:00 2001
From: Sandy Walsh <sandy.walsh@rackspace.com>
Date: Thu, 24 Feb 2011 14:32:25 -0800
Subject: service capabilities test

---
 nova/scheduler/driver.py  |  7 +++++++
 nova/scheduler/manager.py |  3 ++-
 nova/tests/test_zones.py  | 28 ++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/nova/scheduler/driver.py b/nova/scheduler/driver.py
index 66e46c1b9..317a039cc 100644
--- a/nova/scheduler/driver.py
+++ b/nova/scheduler/driver.py
@@ -45,6 +45,13 @@ class WillNotSchedule(exception.Error):
 class Scheduler(object):
     """The base class that all Scheduler clases should inherit from."""
 
+    def __init__(self):
+        self.zone_manager = None
+
+    def set_zone_manager(self, zone_manager):
+        """Called by the Scheduler Service to supply a ZoneManager."""
+        self.zone_manager = zone_manager
+
     @staticmethod
     def service_is_up(service):
         """Check whether a service is up based on last heartbeat."""
diff --git a/nova/scheduler/manager.py b/nova/scheduler/manager.py
index 1bda77d89..d3d338943 100644
--- a/nova/scheduler/manager.py
+++ b/nova/scheduler/manager.py
@@ -41,10 +41,11 @@ flags.DEFINE_string('scheduler_driver',
 class SchedulerManager(manager.Manager):
     """Chooses a host to run instances on."""
     def __init__(self, scheduler_driver=None, *args, **kwargs):
+        self.zone_manager = zone_manager.ZoneManager()
         if not scheduler_driver:
             scheduler_driver = FLAGS.scheduler_driver
         self.driver = utils.import_object(scheduler_driver)
-        self.zone_manager = zone_manager.ZoneManager()
+        self.driver.set_zone_manager(self.zone_manager)
         super(SchedulerManager, self).__init__(*args, **kwargs)
 
     def __getattr__(self, key):
diff --git a/nova/tests/test_zones.py b/nova/tests/test_zones.py
index 5a52a0506..3ca71d5f1 100644
--- a/nova/tests/test_zones.py
+++ b/nova/tests/test_zones.py
@@ -76,6 +76,34 @@ class ZoneManagerTestCase(test.TestCase):
         self.assertEquals(len(zm.zone_states), 1)
         self.assertEquals(zm.zone_states[1].username, 'user1')
 
+    def test_service_capabilities(self):
+        zm = zone_manager.ZoneManager()
+        caps = zm.get_zone_capabilities(self, None)
+        self.assertEquals(caps, {})
+
+        zm.update_service_capabilities("svc1", "host1", dict(a=1, b=2))
+        caps = zm.get_zone_capabilities(self, None)
+        self.assertEquals(caps, dict(svc1_a=(1, 1), svc1_b=(2, 2)))
+
+        zm.update_service_capabilities("svc1", "host1", dict(a=2, b=3))
+        caps = zm.get_zone_capabilities(self, None)
+        self.assertEquals(caps, dict(svc1_a=(2, 2), svc1_b=(3, 3)))
+ 
+        zm.update_service_capabilities("svc1", "host2", dict(a=20, b=30))
+        caps = zm.get_zone_capabilities(self, None)
+        self.assertEquals(caps, dict(svc1_a=(2, 20), svc1_b=(3, 30)))
+  
+        zm.update_service_capabilities("svc10", "host1", dict(a=99, b=99))
+        caps = zm.get_zone_capabilities(self, None)
+        self.assertEquals(caps, dict(svc1_a=(2, 20), svc1_b=(3, 30),
+                                        svc10_a=(99, 99), svc10_b=(99, 99)))
+   
+        zm.update_service_capabilities("svc1", "host3", dict(c=5))
+        caps = zm.get_zone_capabilities(self, None)
+        self.assertEquals(caps, dict(svc1_a=(2, 20), svc1_b=(3, 30),
+                                     svc1_c=(5, 5), svc10_a=(99, 99),
+                                     svc10_b=(99, 99)))
+ 
     def test_refresh_from_db_replace_existing(self):
         zm = zone_manager.ZoneManager()
         zone_state = zone_manager.ZoneState()
-- 
cgit