vm: add Node.flush

closes #60

circle/vm/models/node.py

@@ -13,7 +13,7 @@ from taggit.managers import TaggableManager
 
 
 from common.models import method_cache, WorkerNotFound
 from common.models import method_cache, WorkerNotFound
 from firewall.models import Host
 from firewall.models import Host
-from ..tasks import vm_tasks
+from ..tasks import vm_tasks, local_tasks
 from .common import Trait
 from .common import Trait
 
 
 from .activity import node_activity, NodeActivity
 from .activity import node_activity, NodeActivity
@@ -101,6 +101,21 @@ class Node(TimeStampedModel):
                 self.enabled = False
                 self.enabled = False
                 self.save()
                 self.save()
 
 
+    def flush(self, user=None):
+        """Disable node and move all instances to other ones.
+        """
+        with node_activity('flush', node=self, user=user) as act:
+            self.disable(user, act)
+            for i in self.instance_set.all():
+                with act.sub_activity('migrate_instance_%d' % i.pk):
+                    i.migrate()
+
+    def flush_async(self, user=None):
+        """Execute flush asynchronously.
+        """
+        return local_tasks.flush.apply_async(args=[self, user],
+                                             queue="localhost.man")
+
     def enable(self, user=None):
     def enable(self, user=None):
         ''' Enable the node. '''
         ''' Enable the node. '''
         if self.enabled is not True:
         if self.enabled is not True:

circle/vm/tasks/local_tasks.py

@@ -57,3 +57,8 @@ def reboot(instance, user):
 @celery.task
 @celery.task
 def migrate(instance, to_node,  user):
 def migrate(instance, to_node,  user):
     instance.migrate(to_node, task_uuid=migrate.request.id, user=user)
     instance.migrate(to_node, task_uuid=migrate.request.id, user=user)
+
+
+@celery.task
+def flush(node, user):
+    node.migrate(task_uuid=flush.request.id, user=user)

circle/vm/tests/test_models.py

+from django.contrib.auth.models import User
 from django.test import TestCase
 from django.test import TestCase
 from mock import Mock, MagicMock, patch, call
 from mock import Mock, MagicMock, patch, call
 
 
@@ -72,7 +73,6 @@ class InterfaceTestCase(TestCase):
 
 
     def test_interface_create(self):
     def test_interface_create(self):
         from firewall.models import Vlan, Domain
         from firewall.models import Vlan, Domain
-        from django.contrib.auth.models import User
         owner = User()
         owner = User()
         owner.save()
         owner.save()
         i = Instance(id=10, owner=owner, access_method='rdp')
         i = Instance(id=10, owner=owner, access_method='rdp')
@@ -106,3 +106,67 @@ class LeaseTestCase(TestCase):
 
 
         l.suspend_interval = None
         l.suspend_interval = None
         assert "never" in unicode(l)
         assert "never" in unicode(l)
+
+
+class NodeTestCase(TestCase):
+
+    def test_disable_enabled(self):
+        node = MagicMock(spec=Node, enabled=True)
+        with patch('vm.models.node.node_activity') as nac:
+            na = MagicMock()
+            nac.return_value = na
+            na.__enter__.return_value = MagicMock()
+            Node.disable(node)
+        self.assertFalse(node.enabled)
+        node.save.assert_called_once()
+        na.assert_called()
+
+    def test_disable_disabled(self):
+        node = MagicMock(spec=Node, enabled=False)
+        with patch('vm.models.node.node_activity') as nac:
+            na = MagicMock()
+            na.__enter__.side_effect = AssertionError
+            nac.return_value = na
+            Node.disable(node)
+        self.assertFalse(node.enabled)
+
+    def test_disable_enabled_sub(self):
+        node = MagicMock(spec=Node, enabled=True)
+        act = MagicMock()
+        subact = MagicMock()
+        act.sub_activity.return_value = subact
+        Node.disable(node, base_activity=act)
+        self.assertFalse(node.enabled)
+        subact.__enter__.assert_called()
+
+    def test_flush(self):
+        node = MagicMock(spec=Node, enabled=True)
+        user = MagicMock(spec=User)
+        insts = [MagicMock(spec=Instance), MagicMock(spec=Instance)]
+
+        with patch('vm.models.node.node_activity') as na:
+            act = na.return_value.__enter__.return_value = MagicMock()
+            node.instance_set.all.return_value = insts
+
+            Node.flush(node, user)
+
+            na.__enter__.assert_called()
+            node.disable.assert_called_with(user, act)
+            for i in insts:
+                i.migrate.assert_called()
+
+    def test_flush_disabled_wo_user(self):
+        node = MagicMock(spec=Node, enabled=False)
+        insts = [MagicMock(spec=Instance), MagicMock(spec=Instance)]
+
+        with patch('vm.models.node.node_activity') as na:
+            act = na.return_value.__enter__.return_value = MagicMock()
+            node.instance_set.all.return_value = insts
+
+            Node.flush(node)
+
+            node.disable.assert_called_with(None, act)
+            # ^ should be called, but real method no-ops if disabled
+            na.__enter__.assert_called()
+            for i in insts:
+                i.migrate.assert_called()