1#!/usr/bin/env python3
2#
3#   Copyright 2018 - 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
17import os
18import shutil
19import tempfile
20
21from acts import logger
22from acts.libs.proc import job
23
24_UICD_JAR_CMD = 'java -jar %s/uicd-commandline.jar'
25_UNZIP_CMD = 'tar -xzf %s -C %s'
26
27
28class UicdError(Exception):
29    """Raised for exceptions that occur in UIConductor-related tasks"""
30
31
32class UicdCli(object):
33    """Provides an interface for running UIConductor (Uicd) workflows under its
34    CLI.
35
36    This class does not handle workflow creation, which requires the Uicd
37    frontend.
38    """
39    def __init__(self, uicd_zip, workflow_paths, log_path=None):
40        """Creates a UicdCli object. Extracts the required uicd-cli binaries.
41
42        Args:
43            uicd_zip: The path to uicd_cli.tar.gz
44            workflow_paths: List of paths to uicd workflows and/or directories
45                containing them.
46            log_path: Directory for storing logs generated by Uicd.
47        """
48        # This is done so unit tests can cache the mocked shutil.rmtree value
49        # and call it on __del__ when the patch has been lifted.
50        self._rm_tmpdir = shutil.rmtree
51
52        self._uicd_zip = uicd_zip[0] if isinstance(uicd_zip, list) else uicd_zip
53        self._uicd_path = tempfile.mkdtemp(prefix='uicd')
54        self._log_path = log_path
55        if self._log_path:
56            os.makedirs(self._log_path, exist_ok=True)
57        self._log = logger.create_tagged_trace_logger(tag='Uicd')
58        self._set_workflows(workflow_paths)
59        self._setup_cli()
60
61    def _set_workflows(self, workflow_paths):
62        """Set up a dictionary that maps workflow name to its file location.
63        This allows the user to specify workflows to run without having to
64        provide the full path.
65
66        Args:
67            workflow_paths: List of paths to uicd workflows and/or directories
68                containing them.
69
70        Raises:
71            UicdError if two or more Uicd workflows share the same file name
72        """
73        if isinstance(workflow_paths, str):
74            workflow_paths = [workflow_paths]
75
76        # get a list of workflow files from specified paths
77        def _raise(e):
78            raise e
79        workflow_files = []
80        for path in workflow_paths:
81            if os.path.isfile(path):
82                workflow_files.append(path)
83            else:
84                for (root, _, files) in os.walk(path, onerror=_raise):
85                    for file in files:
86                        workflow_files.append(os.path.join(root, file))
87
88        # populate the dictionary
89        self._workflows = {}
90        for path in workflow_files:
91            workflow_name = os.path.basename(path)
92            if workflow_name in self._workflows.keys():
93                raise UicdError('Uicd workflows may not share the same name.')
94            self._workflows[workflow_name] = path
95
96    def _setup_cli(self):
97        """Extract tar from uicd_zip and place unzipped files in uicd_path.
98
99        Raises:
100            Exception if the extraction fails.
101        """
102        self._log.debug('Extracting uicd-cli binaries from %s' % self._uicd_zip)
103        unzip_cmd = _UNZIP_CMD % (self._uicd_zip, self._uicd_path)
104        try:
105            job.run(unzip_cmd.split())
106        except job.Error:
107            self._log.exception('Failed to extract uicd-cli binaries.')
108            raise
109
110    def run(self, serial, workflows, timeout=120):
111        """Run specified workflows on the UIConductor CLI.
112
113        Args:
114            serial: Device serial
115            workflows: List or str of workflows to run.
116            timeout: Number seconds to wait for command to finish.
117        """
118        base_cmd = _UICD_JAR_CMD % self._uicd_path
119        if isinstance(workflows, str):
120            workflows = [workflows]
121        for workflow_name in workflows:
122            self._log.info('Running workflow "%s"' % workflow_name)
123            if workflow_name in self._workflows:
124                args = '-d %s -i %s' % (serial, self._workflows[workflow_name])
125            else:
126                self._log.error(
127                    'The workflow "%s" does not exist.' % workflow_name)
128                continue
129            if self._log_path:
130                args = '%s -o %s' % (args, self._log_path)
131            cmd = '%s %s' % (base_cmd, args)
132            try:
133                result = job.run(cmd.split(), timeout=timeout)
134            except job.Error:
135                self._log.exception(
136                    'Failed to run workflow "%s"' % workflow_name)
137                continue
138            if result.stdout:
139                stdout_split = result.stdout.splitlines()
140                if len(stdout_split) > 2:
141                    self._log.debug('Uicd logs stored at %s' % stdout_split[2])
142
143    def __del__(self):
144        """Delete the temp directory to Uicd CLI binaries upon ACTS exit."""
145        self._rm_tmpdir(self._uicd_path)
146