| #!/usr/bin/env python3 |
| |
| # Copyright (C) 2024 The Android Open Source Project |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """CLI uploader for Mobly test results to Resultstore.""" |
| |
| import argparse |
| import dataclasses |
| import datetime |
| from importlib import resources |
| import logging |
| import mimetypes |
| import pathlib |
| import platform |
| import shutil |
| import tempfile |
| import warnings |
| from xml.etree import ElementTree |
| |
| import google.auth |
| from google.cloud import storage |
| from googleapiclient import discovery |
| |
| import mobly_result_converter |
| import resultstore_client |
| |
| with warnings.catch_warnings(): |
| warnings.simplefilter('ignore') |
| from google.cloud.storage import transfer_manager |
| |
| logging.getLogger('googleapiclient').setLevel(logging.WARNING) |
| |
| _RESULTSTORE_SERVICE_NAME = 'resultstore' |
| _API_VERSION = 'v2' |
| _DISCOVERY_SERVICE_URL = ( |
| 'https://{api}.googleapis.com/$discovery/rest?version={apiVersion}' |
| ) |
| _TEST_XML = 'test.xml' |
| _TEST_LOGS = 'test.log' |
| _UNDECLARED_OUTPUTS = 'undeclared_outputs/' |
| |
| _TEST_SUMMARY_YAML = 'test_summary.yaml' |
| _TEST_LOG_INFO = 'test_log.INFO' |
| |
| _SUITE_NAME = 'suite_name' |
| _RUN_IDENTIFIER = 'run_identifier' |
| |
| _GCS_BASE_LINK = 'https://console.cloud.google.com/storage/browser' |
| |
| _GCS_UPLOAD_INSTRUCTIONS = ( |
| '\nAutomatic upload to GCS failed.\n' |
| 'Please follow the steps below to manually upload files:\n' |
| f'\t1. Follow the link {_GCS_BASE_LINK}/%s.\n' |
| '\t2. Click "UPLOAD FOLDER".\n' |
| '\t3. Select the directory "%s" to upload.' |
| ) |
| |
| _ResultstoreTreeTags = mobly_result_converter.ResultstoreTreeTags |
| _ResultstoreTreeAttributes = mobly_result_converter.ResultstoreTreeAttributes |
| |
| _Status = resultstore_client.Status |
| |
| |
| @dataclasses.dataclass() |
| class _TestResultInfo: |
| """Info from the parsed test summary used for the Resultstore invocation.""" |
| |
| # Aggregate status of the overall test run. |
| status: _Status = _Status.UNKNOWN |
| # Target ID for the test. |
| target_id: str | None = None |
| |
| |
| def _convert_results( |
| mobly_dir: pathlib.Path, dest_dir: pathlib.Path) -> _TestResultInfo: |
| """Converts Mobly test results into Resultstore artifacts.""" |
| test_result_info = _TestResultInfo() |
| logging.info('Converting raw Mobly logs into Resultstore artifacts...') |
| # Generate the test.xml |
| mobly_yaml_path = mobly_dir.joinpath(_TEST_SUMMARY_YAML) |
| if mobly_yaml_path.is_file(): |
| test_xml = mobly_result_converter.convert(mobly_yaml_path, mobly_dir) |
| ElementTree.indent(test_xml) |
| test_xml.write( |
| str(dest_dir.joinpath(_TEST_XML)), |
| encoding='utf-8', |
| xml_declaration=True, |
| ) |
| test_result_info = _get_test_result_info_from_test_xml(test_xml) |
| |
| # Copy test_log.INFO to test.log |
| test_log_info = mobly_dir.joinpath(_TEST_LOG_INFO) |
| if test_log_info.is_file(): |
| shutil.copyfile(test_log_info, dest_dir.joinpath(_TEST_LOGS)) |
| |
| # Copy directory to undeclared_outputs/ |
| shutil.copytree( |
| mobly_dir, |
| dest_dir.joinpath(_UNDECLARED_OUTPUTS), |
| dirs_exist_ok=True, |
| ) |
| return test_result_info |
| |
| |
| def _get_test_result_info_from_test_xml( |
| test_xml: ElementTree.ElementTree, |
| ) -> _TestResultInfo: |
| """Parses a test_xml element into a _TestResultInfo.""" |
| test_result_info = _TestResultInfo() |
| mobly_suite_element = test_xml.getroot().find( |
| f'./{_ResultstoreTreeTags.TESTSUITE.value}' |
| ) |
| if mobly_suite_element is None: |
| return test_result_info |
| # Set aggregate test status |
| test_result_info.status = _Status.PASSED |
| test_class_elements = mobly_suite_element.findall( |
| f'./{_ResultstoreTreeTags.TESTSUITE.value}') |
| failures = int( |
| mobly_suite_element.get(_ResultstoreTreeAttributes.FAILURES.value) |
| ) |
| errors = int( |
| mobly_suite_element.get(_ResultstoreTreeAttributes.ERRORS.value)) |
| if failures or errors: |
| test_result_info.status = _Status.FAILED |
| else: |
| all_skipped = all([test_case_element.get( |
| _ResultstoreTreeAttributes.RESULT.value) == 'skipped' for |
| test_class_element in test_class_elements for |
| test_case_element in test_class_element.findall( |
| f'./{_ResultstoreTreeTags.TESTCASE.value}')]) |
| if all_skipped: |
| test_result_info.status = _Status.SKIPPED |
| |
| # Set target ID based on test class names, suite name, and custom run |
| # identifier. |
| suite_name_value = None |
| run_identifier_value = None |
| properties_element = mobly_suite_element.find( |
| f'./{_ResultstoreTreeTags.PROPERTIES.value}' |
| ) |
| if properties_element is not None: |
| suite_name = properties_element.find( |
| f'./{_ResultstoreTreeTags.PROPERTY.value}' |
| f'[@{_ResultstoreTreeAttributes.NAME.value}="{_SUITE_NAME}"]' |
| ) |
| if suite_name is not None: |
| suite_name_value = suite_name.get( |
| _ResultstoreTreeAttributes.VALUE.value |
| ) |
| run_identifier = properties_element.find( |
| f'./{_ResultstoreTreeTags.PROPERTY.value}' |
| f'[@{_ResultstoreTreeAttributes.NAME.value}="{_RUN_IDENTIFIER}"]' |
| ) |
| if run_identifier is not None: |
| run_identifier_value = run_identifier.get( |
| _ResultstoreTreeAttributes.VALUE.value |
| ) |
| if suite_name_value: |
| target_id = suite_name_value |
| else: |
| test_class_names = [ |
| test_class_element.get(_ResultstoreTreeAttributes.NAME.value) |
| for test_class_element in test_class_elements |
| ] |
| target_id = '+'.join(test_class_names) |
| if run_identifier_value: |
| target_id = f'{target_id} {run_identifier_value}' |
| |
| test_result_info.target_id = target_id |
| return test_result_info |
| |
| |
| def _upload_dir_to_gcs( |
| src_dir: pathlib.Path, gcs_bucket: str, gcs_dir: str |
| ) -> list[str]: |
| """Uploads the given directory to a GCS bucket.""" |
| # Set correct MIME types for certain text-format files. |
| with resources.as_file( |
| resources.files('data').joinpath('mime.types')) as path: |
| mimetypes.init([path]) |
| |
| bucket_obj = storage.Client().bucket(gcs_bucket) |
| |
| glob = src_dir.rglob('*') |
| file_paths = [ |
| str(path.relative_to(src_dir).as_posix()) |
| for path in glob |
| if path.is_file() |
| ] |
| |
| logging.info( |
| 'Uploading %s files from %s to Cloud Storage directory %s/%s...', |
| len(file_paths), |
| str(src_dir), |
| gcs_bucket, |
| gcs_dir, |
| ) |
| # Ensure that the destination directory has a trailing '/'. |
| blob_name_prefix = gcs_dir |
| if blob_name_prefix and not blob_name_prefix.endswith('/'): |
| blob_name_prefix += '/' |
| |
| # If running on Windows, disable multiprocessing for upload. |
| worker_type = ( |
| transfer_manager.THREAD |
| if platform.system() == 'Windows' |
| else transfer_manager.PROCESS |
| ) |
| results = transfer_manager.upload_many_from_filenames( |
| bucket_obj, |
| file_paths, |
| source_directory=str(src_dir), |
| blob_name_prefix=blob_name_prefix, |
| worker_type=worker_type, |
| ) |
| |
| success_paths = [] |
| for file_name, result in zip(file_paths, results): |
| if isinstance(result, Exception): |
| logging.warning('Failed to upload %s. Error: %s', file_name, result) |
| else: |
| logging.debug('Uploaded %s.', file_name) |
| success_paths.append(file_name) |
| |
| # If all files fail to upload, something wrong happened with the GCS client. |
| # Prompt the user to manually upload the files instead. |
| if file_paths and not success_paths: |
| _prompt_user_upload(src_dir, gcs_bucket) |
| success_paths = file_paths |
| |
| return success_paths |
| |
| |
| def _prompt_user_upload(src_dir: pathlib.Path, gcs_bucket: str) -> None: |
| """Prompts the user to manually upload files to GCS.""" |
| print(_GCS_UPLOAD_INSTRUCTIONS % (gcs_bucket, src_dir)) |
| while True: |
| resp = input( |
| 'Once you see the message "# files successfully uploaded", ' |
| 'enter "Y" or "yes" to continue:') |
| if resp.lower() in ('y', 'yes'): |
| break |
| |
| |
| def _upload_to_resultstore( |
| gcs_bucket: str, |
| gcs_dir: str, |
| file_paths: list[str], |
| status: _Status, |
| target_id: str | None, |
| ) -> None: |
| """Uploads test results to Resultstore.""" |
| logging.info('Generating Resultstore link...') |
| service = discovery.build( |
| _RESULTSTORE_SERVICE_NAME, |
| _API_VERSION, |
| discoveryServiceUrl=_DISCOVERY_SERVICE_URL, |
| ) |
| creds, project_id = google.auth.default() |
| client = resultstore_client.ResultstoreClient(service, creds, project_id) |
| client.create_invocation() |
| client.create_default_configuration() |
| client.create_target(target_id) |
| client.create_configured_target() |
| client.create_action(f'gs://{gcs_bucket}/{gcs_dir}', file_paths) |
| client.set_status(status) |
| client.merge_configured_target() |
| client.finalize_configured_target() |
| client.merge_target() |
| client.finalize_target() |
| client.merge_invocation() |
| client.finalize_invocation() |
| |
| |
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| '-v', '--verbose', action='store_true', help='Enable debug logs.' |
| ) |
| parser.add_argument( |
| 'mobly_dir', |
| help='Directory on host where Mobly results are stored.', |
| ) |
| parser.add_argument( |
| '--gcs_bucket', |
| help='Bucket in GCS where test artifacts are uploaded. If unspecified, ' |
| 'use the current GCP project name as the bucket name.', |
| ) |
| parser.add_argument( |
| '--gcs_dir', |
| help=( |
| 'Directory to save test artifacts in GCS. Specify empty string to ' |
| 'store the files in the bucket root. If unspecified, use the ' |
| 'current timestamp as the GCS directory name.' |
| ), |
| ) |
| parser.add_argument( |
| '--test_title', |
| help='Custom test title to display in the result UI.' |
| ) |
| |
| args = parser.parse_args() |
| logging.basicConfig( |
| format='%(levelname)s: %(message)s', |
| level=(logging.DEBUG if args.verbose else logging.INFO) |
| ) |
| _, project_id = google.auth.default() |
| gcs_bucket = project_id if args.gcs_bucket is None else args.gcs_bucket |
| gcs_dir = ( |
| datetime.datetime.now().strftime('%Y%m%d-%H%M%S') |
| if args.gcs_dir is None |
| else args.gcs_dir |
| ) |
| with tempfile.TemporaryDirectory() as tmp: |
| converted_dir = pathlib.Path(tmp).joinpath(gcs_dir) |
| converted_dir.mkdir() |
| mobly_dir = pathlib.Path(args.mobly_dir).absolute().expanduser() |
| test_result_info = _convert_results(mobly_dir, converted_dir) |
| gcs_files = _upload_dir_to_gcs( |
| converted_dir, gcs_bucket, gcs_dir) |
| _upload_to_resultstore( |
| gcs_bucket, |
| gcs_dir, |
| gcs_files, |
| test_result_info.status, |
| args.test_title or test_result_info.target_id, |
| ) |
| |
| |
| if __name__ == '__main__': |
| main() |
