1#!/usr/bin/env python3
2# Copyright 2023, The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Atest start_avd functions."""
17
18from __future__ import print_function
19
20import argparse
21import json
22import logging
23import os
24from pathlib import Path
25import re
26import subprocess
27import time
28
29from atest import atest_utils as au
30from atest import atest_utils
31from atest import constants
32from atest.atest_enum import DetectType, ExitCode
33from atest.metrics import metrics
34
35ACLOUD_DURATION = 'duration'
36ACLOUD_REPORT_FILE_RE = re.compile(
37    r'.*--report[_-]file(=|\s+)(?P<report_file>[\w/.]+)'
38)
39
40
41def get_report_file(results_dir, acloud_args):
42  """Get the acloud report file path.
43
44  This method can parse either string:
45      --acloud-create '--report-file=/tmp/acloud.json'
46      --acloud-create '--report-file /tmp/acloud.json'
47  and return '/tmp/acloud.json' as the report file. Otherwise returning the
48  default path(/tmp/atest_result/<hashed_dir>/acloud_status.json).
49
50  Args:
51      results_dir: string of directory to store atest results.
52      acloud_args: string of acloud create.
53
54  Returns:
55      A string path of acloud report file.
56  """
57  match = ACLOUD_REPORT_FILE_RE.match(acloud_args)
58  if match:
59    return match.group('report_file')
60  return os.path.join(results_dir, 'acloud_status.json')
61
62
63def acloud_create(report_file, args, no_metrics_notice=True):
64  """Method which runs acloud create with specified args in background.
65
66  Args:
67      report_file: A path string of acloud report file.
68      args: A string of arguments.
69      no_metrics_notice: Boolean whether sending data to metrics or not.
70  """
71  notice = constants.NO_METRICS_ARG if no_metrics_notice else ''
72  match = ACLOUD_REPORT_FILE_RE.match(args)
73  report_file_arg = f'--report-file={report_file}' if not match else ''
74
75  # (b/161759557) Assume yes for acloud create to streamline atest flow.
76  acloud_cmd = (
77      'acloud create -y {ACLOUD_ARGS} {REPORT_FILE_ARG} {METRICS_NOTICE} '
78  ).format(
79      ACLOUD_ARGS=args, REPORT_FILE_ARG=report_file_arg, METRICS_NOTICE=notice
80  )
81  au.colorful_print('\nCreating AVD via acloud...', constants.CYAN)
82  logging.debug('Executing: %s', acloud_cmd)
83  start = time.time()
84  proc = subprocess.Popen(acloud_cmd, shell=True)
85  proc.communicate()
86  acloud_duration = time.time() - start
87  atest_utils.print_and_log_info('"acloud create" process has completed.')
88  # Insert acloud create duration into the report file.
89  result = au.load_json_safely(report_file)
90  if result:
91    result[ACLOUD_DURATION] = acloud_duration
92    try:
93      with open(report_file, 'w+') as _wfile:
94        _wfile.write(json.dumps(result))
95    except OSError as e:
96      atest_utils.print_and_log_error(
97          'Failed dumping duration to the report file: %s', e
98      )
99
100
101def acloud_create_validator(results_dir: str, args: argparse.ArgumentParser):
102  """Check lunch'd target before running 'acloud create'.
103
104  Args:
105      results_dir: A string of the results directory.
106      args: An argparse.Namespace object.
107
108  Returns:
109      If the target is valid:
110          A tuple of (multiprocessing.Process,
111                      report_file path)
112      else:
113          A tuple of (None, None)
114  """
115  target = os.getenv('TARGET_PRODUCT')
116  if not re.match(r'^(aosp_|)cf_.*', target):
117    au.colorful_print(
118        f'{target} is not in cuttlefish family; will not create any AVD.'
119        'Please lunch target which belongs to cuttlefish.',
120        constants.RED,
121    )
122    return None, None
123  if args.start_avd:
124    args.acloud_create = []
125  acloud_args = ' '.join(args.acloud_create)
126  report_file = get_report_file(results_dir, acloud_args)
127  acloud_proc = au.run_multi_proc(
128      func=acloud_create, args=[report_file, acloud_args, args.no_metrics]
129  )
130  return acloud_proc, report_file
131
132
133def probe_acloud_status(report_file, find_build_duration):
134  """Method which probes the 'acloud create' result status.
135
136  If the report file exists and the status is 'SUCCESS', then the creation is
137  successful.
138
139  Args:
140      report_file: A path string of acloud report file.
141      find_build_duration: A float of seconds.
142
143  Returns:
144      0: success.
145      8: acloud creation failure.
146      9: invalid acloud create arguments.
147  """
148  # 1. Created but the status is not 'SUCCESS'
149  if Path(report_file).exists():
150    result = au.load_json_safely(report_file)
151    if not result:
152      return ExitCode.AVD_CREATE_FAILURE
153    if result.get('status') == 'SUCCESS':
154      atest_utils.print_and_log_info('acloud create successfully!')
155      # Always fetch the adb of the first created AVD.
156      adb_port = result.get('data').get('devices')[0].get('adb_port')
157      is_remote_instance = result.get('command') == 'create_cf'
158      adb_ip = '127.0.0.1' if is_remote_instance else '0.0.0.0'
159      os.environ[constants.ANDROID_SERIAL] = f'{adb_ip}:{adb_port}'
160
161      acloud_duration = get_acloud_duration(report_file)
162      if find_build_duration - acloud_duration >= 0:
163        # find+build took longer, saved acloud create time.
164        logging.debug('Saved acloud create time: %ss.', acloud_duration)
165        metrics.LocalDetectEvent(
166            detect_type=DetectType.ACLOUD_CREATE, result=round(acloud_duration)
167        )
168      else:
169        # acloud create took longer, saved find+build time.
170        logging.debug('Saved Find and Build time: %ss.', find_build_duration)
171        metrics.LocalDetectEvent(
172            detect_type=DetectType.FIND_BUILD, result=round(find_build_duration)
173        )
174      return ExitCode.SUCCESS
175    au.colorful_print(
176        'acloud create failed. Please check\n{}\nfor detail'.format(
177            report_file
178        ),
179        constants.RED,
180    )
181    return ExitCode.AVD_CREATE_FAILURE
182
183  # 2. Failed to create because of invalid acloud arguments.
184  msg = 'Invalid acloud arguments found!'
185  au.colorful_print(msg, constants.RED)
186  logging.debug(msg)
187  return ExitCode.AVD_INVALID_ARGS
188
189
190def get_acloud_duration(report_file):
191  """Method which gets the duration of 'acloud create' from a report file.
192
193  Args:
194      report_file: A path string of acloud report file.
195
196  Returns:
197      An float of seconds which acloud create takes.
198  """
199  content = au.load_json_safely(report_file)
200  if not content:
201    return 0
202  return content.get(ACLOUD_DURATION, 0)
203