[autotest] Forward RPCs that update host labels to shards
Other than label_add_hosts and label_remove_hosts,
we have 2 more RPCs that updates host labels.
They are host_add_labels and host_remove_labels.
These RPCs should be forwarded to shards.
BUG=chromium:497887
TEST=puppylab. test_host13 is a host in a shard.
>>> import common
>>> from autotest_lib.server import frontend
>>> afe = frontend.AFE()
>>> h = afe.get_hosts(['test_host13'])[0]
>>> h.add_labels(['l1','l3'])
Adding labels ['l1', 'l3'] to host test_host13
>>> h.remove_labels(['l1'])
DEPLOY=apache
Change-Id: Ic37c026efa009eff3aa5276cbbef8c3e1d62e508
Reviewed-on: https://chromium-review.googlesource.com/276211
Reviewed-by: Mungyung Ryu <mkryu@google.com>
Commit-Queue: Mungyung Ryu <mkryu@google.com>
Tested-by: Mungyung Ryu <mkryu@google.com>
diff --git a/frontend/afe/rpc_utils.py b/frontend/afe/rpc_utils.py
index c043ae7..569fc42 100644
--- a/frontend/afe/rpc_utils.py
+++ b/frontend/afe/rpc_utils.py
@@ -7,6 +7,7 @@
__author__ = 'showard@google.com (Steve Howard)'
import datetime
+from functools import wraps
import inspect
import os
import sys
@@ -1259,12 +1260,17 @@
'SHARD', 'global_afe_hostname')
-def route_rpc_to_master(rpc_name, **kwargs):
+def route_rpc_to_master(func):
"""Route RPC to master AFE.
- @param rpc_name: The name of the rpc.
- @param **kwargs: The kwargs for the rpc.
+ @param func: The function to decorate
+ @returns: The function to replace func with.
"""
- master_afe = frontend.AFE(server=get_global_afe_hostname())
- return master_afe.run(rpc_name, **kwargs)
+ @wraps(func)
+ def replacement(**kwargs):
+ if server_utils.is_shard():
+ master_afe = frontend.AFE(server=get_global_afe_hostname())
+ return master_afe.run(func.func_name, **kwargs)
+ return func(**kwargs)
+ return replacement