storage: Added get_uniform_model_type

src/commissaire/storage/__init__.py

@@ -19,6 +19,7 @@
 import logging
 
 from commissaire.errors import CommissaireError
+from commissaire.models import Model
 
 
 class ConfigurationError(CommissaireError):
@@ -111,3 +112,22 @@ def _list(self, model_instance):
         :rtype: list
         """
         raise NotImplementedError('_list must be overriden.')
+
+
+def get_uniform_model_type(list_of_model_instances):
+    """
+    Returns the model type if the models are of the same type.
+
+    :param list_of_model_instances: List of model instances with data to save
+    :type list_of_model_instances: [commissaire.models.Model, ...]
+    :returns: The model unifrom type
+    :rtype: commissaire.models.Model
+    :raises: TypeError
+    """
+    set_of_types = set([type(x) for x in list_of_model_instances])
+    if len(set_of_types) > 1:
+        raise TypeError('Model instances must be of identical type')
+    first_type = set_of_types.pop()
+    if not issubclass(first_type, Model):
+        raise TypeError('Type must be a Model')
+    return first_type

src/commissaire/storage/client.py

@@ -16,6 +16,7 @@
 import commissaire.models as models
 
 from commissaire.bus import RemoteProcedureCallError
+from commissaire.storage import get_uniform_model_type
 
 
 class StorageClient:
@@ -76,10 +77,7 @@ def get_many(self, list_of_model_instances):
         if len(list_of_model_instances) == 0:
             return []
 
-        set_of_types = set([type(x) for x in list_of_model_instances])
-        if len(set_of_types) > 1:
-            raise TypeError('Model instances must be of identical type')
-        model_class = set_of_types.pop()
+        model_class = get_uniform_model_type(list_of_model_instances)
 
         try:
             model_json_data = [x.to_dict() for x in list_of_model_instances]
@@ -144,10 +142,7 @@ def save_many(self, list_of_model_instances):
         if len(list_of_model_instances) == 0:
             return []
 
-        set_of_types = set([type(x) for x in list_of_model_instances])
-        if len(set_of_types) > 1:
-            raise TypeError('Model instances must be of identical type')
-        model_class = set_of_types.pop()
+        model_class = get_uniform_model_type(list_of_model_instances)
 
         try:
             for model_instance in list_of_model_instances:
@@ -204,10 +199,7 @@ def delete_many(self, list_of_model_instances):
         if len(list_of_model_instances) == 0:
             return
 
-        set_of_types = set([type(x) for x in list_of_model_instances])
-        if len(set_of_types) > 1:
-            raise TypeError('Model instances must be of identical type')
-        model_class = set_of_types.pop()
+        model_class = get_uniform_model_type(list_of_model_instances)
 
         try:
             model_json_data = [x.to_dict() for x in list_of_model_instances]

test/test_storage.py

@@ -0,0 +1,62 @@
+# Copyright (C) 2016  Red Hat, Inc
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""
+Tests for the commissaire.storage module.
+"""
+
+import logging
+import json
+
+from unittest import mock
+
+from . import TestCase
+
+from commissaire import storage
+from commissaire.models import Host, Cluster
+
+
+class TestCommissaireStorage_get_uniform_model_type(TestCase):
+    """
+    Tests for the get_uniform_type function.
+    """
+
+    def test_get_uniform_model_type_with_valid_types(self):
+        """
+        Verify get_uniform_model_type works when types are the same.
+        """
+        self.assertEquals(
+            Host,
+            storage.get_uniform_model_type([
+                Host.new(address='127.0.0.1'),
+                Host.new(address='127.0.0.2')]))
+
+    def test_get_uniform_model_type_with_multiple_types(self):
+        """
+        Verify get_uniform_model_type raises when types are not the same.
+        """
+        self.assertRaises(
+            TypeError,
+            storage.get_uniform_model_type,
+            [Host.new(address='127.0.0.1'),
+             Cluster.new(name='test')])
+
+    def test_get_uniform_model_type_with_invalid_types(self):
+        """
+        Verify get_uniform_model_type raises when types are invalid.
+        """
+        self.assertRaises(
+            TypeError,
+            storage.get_uniform_model_type,
+            [object()])