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 18import argparse 19import json 20import logging 21import socket 22import time 23import threading 24import sys 25 26from host_controller import console 27from host_controller import tfc_host_controller 28from host_controller.build import build_provider_pab 29from host_controller.tfc import tfc_client 30from host_controller.vti_interface import vti_endpoint_client 31from host_controller.tradefed import remote_client 32from 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 42def _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 68def _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 93def 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 parser.add_argument("--password", 129 default=None, 130 help="Password string to pass to the prompt " 131 "when running certain command as root previlege.") 132 parser.add_argument("--flash", 133 default=None, 134 help="GCS URL to an img package. Fetches and flashes " 135 "the device(s) given as the '--serial' flag.") 136 args = parser.parse_args() 137 if args.config_file: 138 config_json = json.load(args.config_file) 139 else: 140 config_json = {} 141 config_json["log_level"] = "DEBUG" 142 config_json["hosts"] = [] 143 host_config = {} 144 host_config["cluster_ids"] = ["local-cluster-1", 145 "local-cluster-2"] 146 host_config["lease_interval_sec"] = 30 147 config_json["hosts"].append(host_config) 148 149 env_vars = env_utils.SaveAndClearEnvVars([_ANDROID_BUILD_TOP]) 150 151 root_logger = logging.getLogger() 152 root_logger.setLevel(getattr(logging, config_json["log_level"])) 153 154 if args.vti: 155 vti_endpoint = vti_endpoint_client.VtiEndpointClient(args.vti) 156 else: 157 vti_endpoint = None 158 159 tfc = None 160 if args.use_tfc: 161 if args.config_file: 162 tfc = tfc_client.CreateTfcClient( 163 config_json["tfc_api_root"], 164 config_json["service_key_json_path"], 165 api_name=config_json["tfc_api_name"], 166 api_version=config_json["tfc_api_version"], 167 scopes=config_json["tfc_scopes"]) 168 else: 169 logging.warning("WARN: If --use_tfc is set, --config_file argument " 170 "value must be provided. Starting without TFC.") 171 172 pab = build_provider_pab.BuildProviderPAB() 173 174 hosts = [] 175 for host_config in config_json["hosts"]: 176 cluster_ids = host_config["cluster_ids"] 177 # If host name is not specified, use local host. 178 hostname = host_config.get("hostname", socket.gethostname()) 179 port = host_config.get("port", remote_client.DEFAULT_PORT) 180 cluster_ids = host_config["cluster_ids"] 181 remote = remote_client.RemoteClient(hostname, port) 182 host = tfc_host_controller.HostController(remote, tfc, hostname, 183 cluster_ids) 184 hosts.append(host) 185 if args.poll: 186 lease_interval_sec = host_config["lease_interval_sec"] 187 host_thread = threading.Thread(target=host.Run, 188 args=(lease_interval_sec,)) 189 host_thread.daemon = True 190 host_thread.start() 191 192 if args.poll: 193 while True: 194 sys.stdin.readline() 195 else: 196 main_console = console.Console(vti_endpoint, tfc, pab, hosts, 197 vti_address=args.vti, 198 password=args.password) 199 if args.vti: 200 main_console.StartJobThreadAndProcessPool() 201 else: 202 logging.warning("vti address is not set. example : " 203 "$ run --vti=<url>") 204 205 try: 206 if args.serial: 207 main_console.SetSerials(args.serial.split(",")) 208 if args.script: 209 if args.loop is None: 210 main_console.ProcessScript(args.script) 211 else: 212 _ScriptLoop(main_console, args.script, args.loop) 213 214 if args.console: 215 main_console.cmdloop() 216 elif args.flash: 217 main_console.FlashImgPackage(args.flash) 218 else: # if not script, the default is console mode. 219 main_console.cmdloop() 220 finally: 221 main_console.TearDown() 222 223 env_utils.RestoreEnvVars(env_vars) 224 225 226if __name__ == "__main__": 227 main() 228