1 #!/usr/bin/env python 2 # 3 # Copyright (C) 2017 The Android Open Source Project 4 # 5 # Licensed under the Apache License, Version 2.0 (the "License"); 6 # you may not use this file except in compliance with the License. 7 # You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 # 17 18 import argparse 19 import json 20 import logging 21 import socket 22 import time 23 import threading 24 import sys 25 26 from host_controller import console 27 from host_controller import tfc_host_controller 28 from host_controller.build import build_provider_pab 29 from host_controller.tfc import tfc_client 30 from host_controller.vti_interface import vti_endpoint_client 31 from host_controller.tradefed import remote_client 32 from vts.utils.python.os import env_utils 33 34 _ANDROID_BUILD_TOP = "ANDROID_BUILD_TOP" 35 _SECONDS_PER_UNIT = { 36 "m": 60, 37 "h": 60 * 60, 38 "d": 60 * 60 * 24 39 } 40 41 42 def _ParseInterval(interval_str): 43 """Parses string to time interval. 44 45 Args: 46 interval_str: string, a floating-point number followed by time unit. 47 48 Returns: 49 float, the interval in seconds. 50 51 Raises: 52 ValueError if the argument format is wrong. 53 """ 54 if not interval_str: 55 raise ValueError("Argument is empty.") 56 57 unit = interval_str[-1] 58 if unit not in _SECONDS_PER_UNIT: 59 raise ValueError("Unknown unit: %s" % unit) 60 61 interval = float(interval_str[:-1]) 62 if interval < 0: 63 raise ValueError("Invalid time interval: %s" % interval) 64 65 return interval * _SECONDS_PER_UNIT[unit] 66 67 68 def _ScriptLoop(hc_console, script_path, loop_interval): 69 """Runs a console script repeatedly. 70 71 Args: 72 hc_console: the host controller console. 73 script_path: string, the path to the script. 74 loop_interval: float or integer, the interval in seconds. 75 """ 76 next_start_time = time.time() 77 while hc_console.ProcessScript(script_path): 78 if loop_interval == 0: 79 continue 80 current_time = time.time() 81 skip_cnt = (current_time - next_start_time) // loop_interval 82 if skip_cnt >= 1: 83 logging.warning("Script execution time is longer than loop " 84 "interval. Skip %d iteration(s).", skip_cnt) 85 next_start_time += (skip_cnt + 1) * loop_interval 86 if next_start_time - current_time >= 0: 87 time.sleep(next_start_time - current_time) 88 else: 89 logging.error("Unexpected timestamps: current=%f, next=%f", 90 current_time, next_start_time) 91 92 93 def main(): 94 """Parses arguments and starts console.""" 95 parser = argparse.ArgumentParser() 96 parser.add_argument("--config-file", 97 default=None, 98 type=argparse.FileType('r'), 99 help="The configuration file in JSON format") 100 parser.add_argument("--poll", action="store_true", 101 help="Disable console and start host controller " 102 "threads polling TFC.") 103 parser.add_argument("--use-tfc", action="store_true", 104 help="Enable TFC (TradeFed Cluster).") 105 parser.add_argument("--vti", 106 default=None, 107 help="The base address of VTI endpoint APIs") 108 parser.add_argument("--script", 109 default=None, 110 help="The path to a script file in .py format") 111 parser.add_argument("--serial", 112 default=None, 113 help="The default serial numbers for flashing and " 114 "testing in the console. Multiple serial numbers " 115 "are separated by comma.") 116 parser.add_argument("--loop", 117 default=None, 118 metavar="INTERVAL", 119 type=_ParseInterval, 120 help="The interval of repeating the script. " 121 "The format is a float followed by unit which is " 122 "one of 'm' (minute), 'h' (hour), and 'd' (day). " 123 "If this option is unspecified, the script will " 124 "be processed once.") 125 parser.add_argument("--console", action="store_true", 126 help="Whether to start a console after processing " 127 "a script.") 128 args = parser.parse_args() 129 if args.config_file: 130 config_json = json.load(args.config_file) 131 else: 132 config_json = {} 133 config_json["log_level"] = "DEBUG" 134 config_json["hosts"] = [] 135 host_config = {} 136 host_config["cluster_ids"] = ["local-cluster-1", 137 "local-cluster-2"] 138 host_config["lease_interval_sec"] = 30 139 config_json["hosts"].append(host_config) 140 141 env_vars = env_utils.SaveAndClearEnvVars([_ANDROID_BUILD_TOP]) 142 143 root_logger = logging.getLogger() 144 root_logger.setLevel(getattr(logging, config_json["log_level"])) 145 146 if args.vti: 147 vti_endpoint = vti_endpoint_client.VtiEndpointClient(args.vti) 148 else: 149 vti_endpoint = None 150 151 tfc = None 152 if args.use_tfc: 153 if args.config_file: 154 tfc = tfc_client.CreateTfcClient( 155 config_json["tfc_api_root"], 156 config_json["service_key_json_path"], 157 api_name=config_json["tfc_api_name"], 158 api_version=config_json["tfc_api_version"], 159 scopes=config_json["tfc_scopes"]) 160 else: 161 print("WARN: If --use_tfc is set, --config_file argument value " 162 "must be provided. Starting without TFC.") 163 164 pab = build_provider_pab.BuildProviderPAB() 165 166 hosts = [] 167 for host_config in config_json["hosts"]: 168 cluster_ids = host_config["cluster_ids"] 169 # If host name is not specified, use local host. 170 hostname = host_config.get("hostname", socket.gethostname()) 171 port = host_config.get("port", remote_client.DEFAULT_PORT) 172 cluster_ids = host_config["cluster_ids"] 173 remote = remote_client.RemoteClient(hostname, port) 174 host = tfc_host_controller.HostController(remote, tfc, hostname, 175 cluster_ids) 176 hosts.append(host) 177 if args.poll: 178 lease_interval_sec = host_config["lease_interval_sec"] 179 host_thread = threading.Thread(target=host.Run, 180 args=(lease_interval_sec,)) 181 host_thread.daemon = True 182 host_thread.start() 183 184 if args.poll: 185 while True: 186 sys.stdin.readline() 187 else: 188 main_console = console.Console(vti_endpoint, tfc, pab, hosts, 189 vti_address=args.vti) 190 main_console.StartJobThreadAndProcessPool() 191 try: 192 if args.serial: 193 main_console.SetSerials(args.serial.split(",")) 194 if args.script: 195 if args.loop is None: 196 main_console.ProcessScript(args.script) 197 else: 198 _ScriptLoop(main_console, args.script, args.loop) 199 200 if args.console: 201 main_console.cmdloop() 202 else: # if not script, the default is console mode. 203 main_console.cmdloop() 204 finally: 205 main_console.TearDown() 206 207 env_utils.RestoreEnvVars(env_vars) 208 209 210 if __name__ == "__main__": 211 main() 212
