autotest: add a remove_boards_from_shard RPC
BUG=chromium:704445
TEST=unittests added, and they pass
Change-Id: I1dbbc68cbb29a341e9d35d1bdb322f1bb398da74
Reviewed-on: https://chromium-review.googlesource.com/467990
Commit-Ready: Aviv Keshet <akeshet@chromium.org>
Tested-by: Aviv Keshet <akeshet@chromium.org>
Reviewed-by: Dan Shi <dshi@google.com>
diff --git a/frontend/afe/rpc_interface.py b/frontend/afe/rpc_interface.py
index e17a72c..2ea3d65 100644
--- a/frontend/afe/rpc_interface.py
+++ b/frontend/afe/rpc_interface.py
@@ -37,6 +37,7 @@
import os
import sys
+from django.db import transaction
from django.db.models import Count
import common
@@ -2063,6 +2064,30 @@
return shard.id
+# Remove board RPCs are rare, so we can afford to make them a bit more
+# expensive (by performing in a transaction) in order to guarantee
+# atomicity.
+# TODO(akeshet): If we ever update to newer version of django, we need to
+# migrate to transaction.atomic instead of commit_on_success
+@transaction.commit_on_success
+def remove_board_from_shard(hostname, label):
+ """Remove board from the given shard.
+ @param hostname: The hostname of the shard to be changed.
+ @param labels: Board label.
+
+ @raises models.Label.DoesNotExist: If the label specified doesn't exist.
+
+ @returns: The id of the changed shard.
+ """
+ shard = models.Shard.objects.get(hostname=hostname)
+ label = models.Label.smart_get(label)
+ if label not in shard.labels.all():
+ raise error.RPCException(
+ 'Cannot remove label from shard that does not belong to it.')
+ shard.labels.remove(label)
+ models.Host.objects.filter(labels__in=[label]).update(shard=None)
+
+
def delete_shard(hostname):
"""Delete a shard and reclaim all resources from it.
diff --git a/frontend/afe/rpc_interface_unittest.py b/frontend/afe/rpc_interface_unittest.py
index c62f42a..7b9fc29 100755
--- a/frontend/afe/rpc_interface_unittest.py
+++ b/frontend/afe/rpc_interface_unittest.py
@@ -1116,13 +1116,8 @@
# shard1.
self.mox.StubOutWithMock(models.Host, '_assign_to_shard_nothing_helper')
def remove_label():
- # A separate RPC call to remove_board_from_shard (not yet
- # implemented) swoops in and removes the label from the shard and
- # the shard from all previously assigned hosts.
- shard1.labels.remove(lumpy_label)
- models.Host.objects.filter(
- labels__in=[lumpy_label]
- ).update(shard=None)
+ rpc_interface.remove_board_from_shard(
+ shard1.hostname, lumpy_label.name)
models.Host._assign_to_shard_nothing_helper().WithSideEffects(
remove_label)
self.mox.ReplayAll()
@@ -1133,6 +1128,30 @@
self.assertEqual(host2.shard, None)
+ def testShardLabelRemovalInvalid(self):
+ """Ensure you cannot remove the wrong label from shard."""
+ shard1, host1, lumpy_label = self._createShardAndHostWithLabel()
+ stumpy_label = models.Label.objects.create(
+ name='board:stumpy', platform=True)
+ with self.assertRaises(error.RPCException):
+ rpc_interface.remove_board_from_shard(
+ shard1.hostname, stumpy_label.name)
+
+
+ def testShardHeartbeatLabelRemoval(self):
+ """Ensure label removal from shard works."""
+ shard1, host1, lumpy_label = self._createShardAndHostWithLabel()
+
+ self.assertEqual(host1.shard, shard1)
+ self.assertItemsEqual(shard1.labels.all(), [lumpy_label])
+ rpc_interface.remove_board_from_shard(
+ shard1.hostname, lumpy_label.name)
+ host1 = models.Host.smart_get(host1.id)
+ shard1 = models.Shard.smart_get(shard1.id)
+ self.assertEqual(host1.shard, None)
+ self.assertItemsEqual(shard1.labels.all(), [])
+
+
def testShardRetrieveJobs(self):
"""Create jobs and retrieve them."""
# should never be returned by heartbeat
