1#!/usr/bin/env python3
2#
3#   Copyright 2019 - 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 shlex
19
20DEFAULT_INSTRUMENTATION_LOG_OUTPUT = 'instrumentation_output.txt'
21
22
23class InstrumentationCommandBuilder(object):
24    """Helper class to build instrumentation commands."""
25
26    def __init__(self):
27        self._manifest_package_name = None
28        self._flags = []
29        self._key_value_params = {}
30        self._runner = None
31        self._nohup = False
32        self._proto_path = None
33        self._nohup_log_path = None
34        self._output_as_proto = False
35
36    def set_manifest_package(self, test_package):
37        self._manifest_package_name = test_package
38
39    def set_runner(self, runner):
40        self._runner = runner
41
42    def add_flag(self, param):
43        self._flags.append(param)
44
45    def remove_flag(self, param):
46        while self._flags.count(param):
47            self._flags.remove(param)
48
49    def add_key_value_param(self, key, value):
50        if isinstance(value, bool):
51            value = str(value).lower()
52        self._key_value_params[key] = str(value)
53
54    def set_proto_path(self, path=None):
55        """Sets a custom path to store result proto. Note that this path will
56        be relative to $EXTERNAL_STORAGE on device. Calling this function
57        automatically enables output as proto.
58
59        Args:
60            path: The $EXTERNAL_STORAGE subdirectory to write the result proto
61            to. If left as None, the default location will be used.
62        """
63        self._output_as_proto = True
64        self._proto_path = path
65
66    def set_output_as_text(self):
67        """This is the default behaviour. It will simply output the
68        instrumentation output to the devices stdout. If the nohup option is
69        enabled the instrumentation output will be redirected to the defined
70        path or its default.
71        """
72        self._output_as_proto = False
73        self._proto_path = None
74
75    def set_nohup(self, log_path=None):
76        """Enables nohup mode. This enables the instrumentation command to
77        continue running after a USB disconnect.
78
79        Args:
80            log_path: Path to store stdout of the process. Default is:
81            $EXTERNAL_STORAGE/instrumentation_output.txt
82        """
83        if log_path is None:
84            log_path = os.path.join('$EXTERNAL_STORAGE',
85                                    DEFAULT_INSTRUMENTATION_LOG_OUTPUT)
86        self._nohup = True
87        self._nohup_log_path = log_path
88
89    def build(self):
90        call = self._instrument_call_with_arguments()
91        call.append('{}/{}'.format(self._manifest_package_name, self._runner))
92        if self._nohup:
93            call = ['nohup'] + call
94            call.append('>>')
95            call.append(self._nohup_log_path)
96            call.append('2>&1')
97        return " ".join(call)
98
99    def _instrument_call_with_arguments(self):
100        errors = []
101        if self._manifest_package_name is None:
102            errors.append('manifest package cannot be none')
103        if self._runner is None:
104            errors.append('instrumentation runner cannot be none')
105        if len(errors) > 0:
106            raise Exception('instrumentation call build errors: {}'
107                            .format(','.join(errors)))
108        call = ['am instrument']
109
110        for flag in self._flags:
111            call.append(flag)
112
113        if self._output_as_proto:
114            call.append('-f')
115        if self._proto_path:
116            call.append(self._proto_path)
117        for key, value in self._key_value_params.items():
118            call.append('-e')
119            call.append(key)
120            call.append(shlex.quote(value))
121        return call
122
123
124class InstrumentationTestCommandBuilder(InstrumentationCommandBuilder):
125
126    def __init__(self):
127        super().__init__()
128        self._packages = []
129        self._classes = []
130
131    @staticmethod
132    def default():
133        """Default instrumentation call builder.
134
135        The flags -w, -r are enabled.
136
137           -w  Forces am instrument to wait until the instrumentation terminates
138           (needed for logging)
139           -r  Outputs results in raw format.
140           https://developer.android.com/studio/test/command-line#AMSyntax
141
142        The default test runner is androidx.test.runner.AndroidJUnitRunner.
143        """
144        builder = InstrumentationTestCommandBuilder()
145        builder.add_flag('-w')
146        builder.add_flag('-r')
147        builder.set_runner('androidx.test.runner.AndroidJUnitRunner')
148        return builder
149
150    CONFLICTING_PARAMS_MESSAGE = ('only a list of classes and test methods or '
151                                  'a list of test packages are allowed.')
152
153    def add_test_package(self, package):
154        if len(self._classes) != 0:
155            raise Exception(self.CONFLICTING_PARAMS_MESSAGE)
156        self._packages.append(package)
157
158    def add_test_method(self, class_name, test_method):
159        if len(self._packages) != 0:
160            raise Exception(self.CONFLICTING_PARAMS_MESSAGE)
161        self._classes.append('{}#{}'.format(class_name, test_method))
162
163    def add_test_class(self, class_name):
164        if len(self._packages) != 0:
165            raise Exception(self.CONFLICTING_PARAMS_MESSAGE)
166        self._classes.append(class_name)
167
168    def build(self):
169        errors = []
170        if len(self._packages) == 0 and len(self._classes) == 0:
171            errors.append('at least one of package, class or test method need '
172                          'to be defined')
173
174        if len(errors) > 0:
175            raise Exception('instrumentation call build errors: {}'
176                            .format(','.join(errors)))
177
178        call = self._instrument_call_with_arguments()
179
180        if len(self._packages) > 0:
181            call.append('-e')
182            call.append('package')
183            call.append(','.join(self._packages))
184        elif len(self._classes) > 0:
185            call.append('-e')
186            call.append('class')
187            call.append(','.join(self._classes))
188
189        call.append('{}/{}'.format(self._manifest_package_name, self._runner))
190        if self._nohup:
191            call = ['nohup'] + call
192            call.append('>>')
193            call.append(self._nohup_log_path)
194            call.append('2>&1')
195        return ' '.join(call)
196