[autotest] Add tests for _allowed_hosts_for_master_job
BUG=None
TEST=Run tests
Change-Id: I844e91746e94467e0ef2e19481fd2626575f2eff
Reviewed-on: https://chromium-review.googlesource.com/420353
Commit-Ready: Allen Li <ayatane@chromium.org>
Tested-by: Allen Li <ayatane@chromium.org>
Reviewed-by: Prathmesh Prabhu <pprabhu@chromium.org>
diff --git a/frontend/afe/rpc_utils_unittest.py b/frontend/afe/rpc_utils_unittest.py
index 5a64f42..d430ad7 100755
--- a/frontend/afe/rpc_utils_unittest.py
+++ b/frontend/afe/rpc_utils_unittest.py
@@ -6,7 +6,7 @@
"""Unit tests for frontend/afe/rpc_utils.py."""
-
+import mock
import unittest
import common
@@ -87,5 +87,42 @@
self.assertEquals(got, {'a': 1, 'b': 2, 'args': (3,)})
+class AllowedHostsForMasterJobTest(unittest.TestCase):
+ """Unit tests for _allowed_hosts_for_master_job()."""
+
+ # pylint: disable=missing-docstring
+
+ @mock.patch.object(rpc_utils, 'bucket_hosts_by_shard', autospec=True)
+ def test_multiple_shards(self, bucket_mock):
+ bucket_mock.return_value = {
+ 'shard1': [],
+ 'shard2': [],
+ }
+ got = rpc_utils._allowed_hosts_for_master_job([])
+ self.assertFalse(got)
+
+ @mock.patch.object(rpc_utils, 'bucket_hosts_by_shard', autospec=True)
+ def test_one_shard_with_less_hosts(self, bucket_mock):
+ bucket_mock.return_value = {
+ 'shard1': [1],
+ }
+ got = rpc_utils._allowed_hosts_for_master_job([1, 2])
+ self.assertFalse(got)
+
+ @mock.patch.object(rpc_utils, 'bucket_hosts_by_shard', autospec=True)
+ def test_one_shard_with_equal_hosts(self, bucket_mock):
+ bucket_mock.return_value = {
+ 'shard1': [1, 2],
+ }
+ got = rpc_utils._allowed_hosts_for_master_job([1, 2])
+ self.assertTrue(got)
+
+ @mock.patch.object(rpc_utils, 'bucket_hosts_by_shard', autospec=True)
+ def test_no_shards(self, bucket_mock):
+ bucket_mock.return_value = {}
+ got = rpc_utils._allowed_hosts_for_master_job([1, 2])
+ self.assertTrue(got)
+
+
if __name__ == '__main__':
unittest.main()