[autotest] Restructure _allowed_hosts_for_master_job() logic

BUG=None
TEST=Unit tests

Change-Id: I4573bf69abbef91cdbb96960865bb8cc0232becf
Reviewed-on: https://chromium-review.googlesource.com/430206
Commit-Ready: Allen Li <ayatane@chromium.org>
Tested-by: Allen Li <ayatane@chromium.org>
Reviewed-by: Allen Li <ayatane@chromium.org>
diff --git a/frontend/afe/rpc_utils.py b/frontend/afe/rpc_utils.py
index 4d808d2..14207ae 100644
--- a/frontend/afe/rpc_utils.py
+++ b/frontend/afe/rpc_utils.py
@@ -1031,16 +1031,21 @@
 
 def _allowed_hosts_for_master_job(host_objects):
     """Check that the hosts are allowed for a job on master."""
-    shard_host_map = bucket_hosts_by_shard(host_objects)
-    num_shards = len(shard_host_map)
     # We disallow the following jobs on master:
     #   num_shards > 1: this is a job spanning across multiple shards.
     #   num_shards == 1 but number of hosts on shard is less
     #   than total number of hosts: this is a job that spans across
     #   one shard and the master.
-    return not (num_shards > 1
-                or (num_shards == 1
-                    and len(shard_host_map.values()[0]) != len(host_objects)))
+    shard_host_map = bucket_hosts_by_shard(host_objects)
+    num_shards = len(shard_host_map)
+    if num_shards > 1:
+        return False
+    if num_shards == 1:
+        hosts_on_shard = shard_host_map.values()[0]
+        assert len(hosts_on_shard) <= len(host_objects)
+        return len(hosts_on_shard) == len(host_objects)
+    else:
+        return True
 
 
 def encode_ascii(control_file):