tests and better driver loading

3 files changed, 127 insertions, 20 deletions

diff --git a/nova/exception.py b/nova/exception.py
index e8444cb14..586ec009b 100644
--- a/nova/exception.py
+++ b/nova/exception.py
@@ -449,6 +449,14 @@ class ZoneNotFound(NotFound):
     message = _("Zone %(zone_id)s could not be found.")
 
 
+class SchedulerQueryDriverNotFound(NotFound):
+    message = _("Scheduler Query Driver %(driver_name)s could not be found.")
+
+
+class BadSchedulerQueryDriver(NotFound):
+    message = _("Invalid Scheduler Query Driver selected.")
+
+
 class InstanceMetadataNotFound(NotFound):
     message = _("Instance %(instance_id)s has no metadata with "
                 "key %(metadata_key)s.")
diff --git a/nova/scheduler/query.py b/nova/scheduler/query.py
index 5379b3d44..4971bc7e4 100644
--- a/nova/scheduler/query.py
+++ b/nova/scheduler/query.py
@@ -15,39 +15,38 @@
 
 """
 Query is a plug-in mechanism for requesting instance resources.
-Three plug-ins are included: PassThru, Flavor & JSON. PassThru just
+Three plug-ins are included: AllHosts, Flavor & JSON. AllHosts just
 returns the full, unfiltered list of hosts. Flavor is a hard coded
 matching mechanism based on flavor criteria and JSON is an ad-hoc
 query grammar.
 """
 
 from nova import exception
+from nova import flags
+from nova import log as logging
 from nova import utils
 
+LOG = logging.getLogger('nova.scheduler.query')
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('default_query_engine',
+                    'nova.scheduler.query.AllHostsQuery',
+                    'Which query engine to use for filtering hosts.')
+
+
 class Query:
     """Base class for query plug-ins."""
 
     def instance_type_to_query(self, instance_type):
         """Convert instance_type into a query for most common use-case."""
-        raise exception.Error(_("Query driver not specified."))
+        raise exception.BadSchedulerQueryDriver()
 
     def filter_hosts(self, zone_manager, query):
         """Return a list of hosts that fulfill the query."""
-        raise exception.Error(_("Query driver not specified."))
-
+        raise exception.BadSchedulerQueryDriver()
 
-def load_driver(driver_name):
-    resource = utils.import_class(driver_name)
-    if type(resource) != types.ClassType:
-        continue
-    cls = resource
-    if not issubclass(cls, Query):
-        raise exception.Error(_("Query driver does not derive "
-                "from nova.scheduler.query.Query."))
-    return cls
 
-
-class PassThruQuery:
+class AllHostsQuery:
     """NOP query plug-in. Returns all hosts in ZoneManager.
     This essentially does what the old Scheduler+Chance used
     to give us."""
@@ -59,7 +58,7 @@ class PassThruQuery:
 
     def filter_hosts(self, zone_manager, query):
         """Return a list of hosts from ZoneManager list."""
-        hosts = zone_manager.service_state.get('compute', {})
+        hosts = zone_manager.service_states.get('compute', {})
         return [(host, capabilities)
                 for host, capabilities in hosts.iteritems()]
 
@@ -73,7 +72,7 @@ class FlavorQuery:
 
     def filter_hosts(self, zone_manager, query):
         """Return a list of hosts that can create instance_type."""
-        hosts = zone_manager.service_state.get('compute', {})
+        hosts = zone_manager.service_states.get('compute', {})
         selected_hosts = []
         instance_type = query
         for host, capabilities in hosts.iteritems():
@@ -148,8 +147,8 @@ class JsonQuery:
         '<=': _less_than_equal,
         '>=': _greater_than_equal,
         'not': _not,
-        'must', _must,
-        'or', _or,
+        'must': _must,
+        'or': _or,
     }
 
     def instance_type_to_query(self, instance_type):
@@ -161,4 +160,15 @@ class JsonQuery:
         return []
 
 
-                              
+# Since the caller may specify which driver to use we need
+# to have an authoritative list of what is permissible. 
+DRIVERS = [AllHostsQuery, FlavorQuery, JsonQuery]
+
+
+def choose_driver(driver_name=None):
+    if not driver_name:
+        driver_name = FLAGS.default_query_engine
+    for driver in DRIVERS:
+        if str(driver) == driver_name:
+            return driver
+    raise exception.SchedulerQueryDriverNotFound(driver_name=driver_name)
diff --git a/nova/tests/test_query.py b/nova/tests/test_query.py
new file mode 100644
index 000000000..0fab91933
--- /dev/null
+++ b/nova/tests/test_query.py
@@ -0,0 +1,89 @@
+# Copyright 2011 OpenStack LLC. 
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+"""
+Tests For Scheduler Query Drivers
+"""
+
+from nova import exception
+from nova import flags
+from nova import test
+from nova.scheduler import query
+
+FLAGS = flags.FLAGS
+
+class FakeZoneManager:
+    pass
+
+class QueryTestCase(test.TestCase):
+    """Test case for query drivers."""
+
+    def _host_caps(self, multiplier):
+        return {'host_name-description':'XenServer %s' % multiplier,
+                'host_hostname':'xs-%s' % multiplier,
+                'host_memory':{'total': 100,
+                             'overhead': 5,
+                             'free': 5 + multiplier * 5,
+                             'free-computed': 5 + multiplier * 5},
+                'host_other-config':{},
+                'host_ip_address':'192.168.1.%d' % (100 + multiplier),
+                'host_cpu_info':{},
+                'disk':{'available': 100 + multiplier * 100,
+                      'total': 1000,
+                      'used': 50 + multiplier * 50},
+                'host_uuid':'xxx-%d' % multiplier,
+                'host_name-label':'xs-%s' % multiplier}
+
+    def setUp(self):
+        self.old_flag = FLAGS.default_query_engine
+        FLAGS.default_query_engine = 'nova.scheduler.query.AllHostsQuery'
+        self.instance_type = dict(name='tiny',
+                memory_mb=500,
+                vcpus=10,
+                local_gb=50,
+                flavorid=1,
+                swap=500,
+                rxtx_quota=30000,
+                rxtx_cap=200)
+
+        hosts = {}
+        for x in xrange(10):
+            hosts['host%s' % x] = self._host_caps(x)
+
+        self.zone_manager = FakeZoneManager()
+        self.zone_manager.service_states = {}
+        self.zone_manager.service_states['compute'] = hosts
+
+    def tearDown(self):
+        FLAGS.default_query_engine = self.old_flag
+
+    def test_choose_driver(self):
+        driver = query.choose_driver()
+        self.assertEquals(str(driver), 'nova.scheduler.query.AllHostsQuery')
+        driver = query.choose_driver('nova.scheduler.query.FlavorQuery')
+        self.assertEquals(str(driver), 'nova.scheduler.query.FlavorQuery')
+        try:
+            query.choose_driver('does not exist')
+            self.fail("Should not find driver")
+        except exception.SchedulerQueryDriverNotFound:
+            pass
+
+    def test_all_host_driver(self):
+        driver = query.AllHostsQuery()
+        cooked = driver.instance_type_to_query(self.instance_type)
+        hosts = driver.filter_hosts(self.zone_manager, cooked)
+        self.assertEquals(10, len(hosts))
+        for host, capabilities in hosts:
+            self.assertTrue(host.startswith('host'))
+