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