Check for multiple platforms on a machine when modifying labels.

Signed-off-by: Steve Howard <showard@google.com>


git-svn-id: svn://test.kernel.org/autotest/trunk@3179 592f7852-d20e-0410-864c-8624ca9c26a4

3 files changed, 52 insertions, 6 deletions

diff --git a/frontend/afe/models.py b/frontend/afe/models.py
index 332a95f5..559b3a78 100644
--- a/frontend/afe/models.py
+++ b/frontend/afe/models.py
@@ -1,7 +1,8 @@
 from datetime import datetime
 from django.db import models as dbmodels, connection
-from frontend.afe import model_logic
-from frontend import settings, thread_local
+import common
+from autotest_lib.frontend.afe import model_logic
+from autotest_lib.frontend import settings, thread_local
 from autotest_lib.client.common_lib import enum, host_protections, global_config
 from autotest_lib.client.common_lib import debug
 
@@ -348,6 +349,22 @@ class Host(model_logic.ModelWithInvalid, dbmodels.Model,
     platform.short_description = 'Platform'
 
 
+    @classmethod
+    def check_no_platform(cls, hosts):
+        Host.objects.populate_relationships(hosts, Label, 'label_list')
+        errors = []
+        for host in hosts:
+            platforms = [label.name for label in host.label_list
+                         if label.platform]
+            if platforms:
+                # do a join, just in case this host has multiple platforms,
+                # we'll be able to see it
+                errors.append('Host %s already has a platform: %s' % (
+                              host.hostname, ', '.join(platforms)))
+        if errors:
+            raise model_logic.ValidationError({'labels': '; '.join(errors)})
+
+
     def is_dead(self):
         return self.status == Host.Status.REPAIR_FAILED
 
diff --git a/frontend/afe/rpc_interface.py b/frontend/afe/rpc_interface.py
index 38951902..6f9e5cc4 100644
--- a/frontend/afe/rpc_interface.py
+++ b/frontend/afe/rpc_interface.py
@@ -30,8 +30,10 @@ See doctests/001_rpc_test.txt for (lots) more examples.
 __author__ = 'showard@google.com (Steve Howard)'
 
 import datetime
-from frontend import thread_local
-from frontend.afe import models, model_logic, control_file, rpc_utils
+import common
+from autotest_lib.frontend import thread_local
+from autotest_lib.frontend.afe import models, model_logic, control_file
+from autotest_lib.frontend.afe import rpc_utils
 from autotest_lib.client.common_lib import global_config
 
 
@@ -53,7 +55,10 @@ def delete_label(id):
 
 def label_add_hosts(id, hosts):
     host_objs = models.Host.smart_get_bulk(hosts)
-    models.Label.smart_get(id).host_set.add(*host_objs)
+    label = models.Label.smart_get(id)
+    if label.platform:
+        models.Host.check_no_platform(host_objs)
+    label.host_set.add(*host_objs)
 
 
 def label_remove_hosts(id, hosts):
@@ -124,7 +129,16 @@ def modify_hosts(host_filter_data, update_data):
 
 def host_add_labels(id, labels):
     labels = models.Label.smart_get_bulk(labels)
-    models.Host.smart_get(id).labels.add(*labels)
+    host = models.Host.smart_get(id)
+
+    platforms = [label.name for label in labels if label.platform]
+    if len(platforms) > 1:
+        raise model_logic.ValidationError(
+            {'labels': 'Adding more than one platform label: %s' %
+                       ', '.join(platforms)})
+    if len(platforms) == 1:
+        models.Host.check_no_platform([host])
+    host.labels.add(*labels)
 
 
 def host_remove_labels(id, labels):
diff --git a/frontend/afe/rpc_interface_unittest.py b/frontend/afe/rpc_interface_unittest.py
index 527c90ae..771cc716 100644
--- a/frontend/afe/rpc_interface_unittest.py
+++ b/frontend/afe/rpc_interface_unittest.py
@@ -6,6 +6,7 @@ from autotest_lib.frontend import setup_django_environment
 from autotest_lib.frontend.afe import frontend_test_utils
 from django.db import connection
 from autotest_lib.frontend.afe import models, rpc_interface, frontend_test_utils
+from autotest_lib.frontend.afe import model_logic
 
 
 _hqe_status = models.HostQueueEntry.Status
@@ -21,6 +22,20 @@ class RpcInterfaceTest(unittest.TestCase,
         self._frontend_common_teardown()
 
 
+    def test_multiple_platforms(self):
+        platform2 = models.Label.objects.create(name='platform2', platform=True)
+        self.assertRaises(model_logic.ValidationError,
+                          rpc_interface. label_add_hosts, 'platform2',
+                          ['host1', 'host2'])
+        self.assertRaises(model_logic.ValidationError,
+                          rpc_interface.host_add_labels, 'host1', ['platform2'])
+        # make sure the platform didn't get added
+        platforms = rpc_interface.get_labels(
+            host__hostname__in=['host1', 'host2'], platform=True)
+        self.assertEquals(len(platforms), 1)
+        self.assertEquals(platforms[0]['name'], 'myplatform')
+
+
     def test_get_jobs_summary(self):
         job = self._create_job(xrange(3))
         entries = list(job.hostqueueentry_set.all())