1# Copyright 2018, The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15ATEST_REL_DIR="tools/asuite/atest"
16
17_atest_completion_ready() {
18    # Not support completion on systems of which the out-of-box bash version
19    # is too old to fully support required built-in commands/functions.
20    # For MacOS users, atest tab completion can be enabled via:
21    #    brew install bash bash-completion@2
22    # and configure code below in ~/.bashrc
23    #    completion_file="/usr/local/etc/profile.d/bash_completion.sh"
24    #    [[ -r "$completion_file" ]] && source "$completion_file"
25    # Open a new terminal, source/lunch and try again.
26    reqs=(compopt _get_comp_words_by_ref __ltrim_colon_completions)
27    for _cmd in "${reqs[@]}"; do
28        if ! type "$_cmd" >/dev/null 2>&1; then
29            return 1
30        fi
31    done
32}
33
34_fetch_testable_modules() {
35    [[ -z $ANDROID_BUILD_TOP ]] && return 0
36    export ATEST_DIR="$ANDROID_BUILD_TOP/$ATEST_REL_DIR"
37    $PYTHON - << END
38import os
39import pickle
40import sys
41
42sys.path.append(os.getenv('ATEST_DIR'))
43import constants
44
45if os.path.isfile(constants.MODULE_INDEX):
46    with open(constants.MODULE_INDEX, 'rb') as cache:
47        try:
48            print("\n".join(pickle.load(cache, encoding="utf-8")))
49        except:
50            print("\n".join(pickle.load(cache)))
51else:
52    print("")
53END
54    unset ATEST_DIR
55}
56
57# This function invoke get_args() and return each item
58# of the list for tab completion candidates.
59_fetch_atest_args() {
60    [[ -z $ANDROID_BUILD_TOP ]] && return 0
61    export ATEST_DIR="$ANDROID_BUILD_TOP/$ATEST_REL_DIR"
62    $PYTHON - << END
63import os
64import sys
65
66atest_dir = os.path.join(os.getenv('ATEST_DIR'))
67sys.path.append(atest_dir)
68
69import atest_arg_parser
70
71parser = atest_arg_parser.AtestArgParser()
72parser.add_atest_args()
73print("\n".join(parser.get_args()))
74END
75    unset ATEST_DIR
76}
77
78# This function returns devices recognised by adb.
79_fetch_adb_devices() {
80    while read dev; do echo $dev | awk '{print $1}'; done < <(adb devices | egrep -v "^List|^$"||true)
81}
82
83# This function returns all paths contain TEST_MAPPING.
84_fetch_test_mapping_files() {
85    [[ -z $ANDROID_BUILD_TOP ]] && return 0
86    find -maxdepth 5 -type f -name TEST_MAPPING |sed 's/^.\///g'| xargs dirname 2>/dev/null
87}
88
89# The main tab completion function.
90_atest() {
91    local cur prev
92    COMPREPLY=()
93    cur="${COMP_WORDS[COMP_CWORD]}"
94    prev="${COMP_WORDS[COMP_CWORD-1]}"
95    _get_comp_words_by_ref -n : cur prev || true
96
97    case "$cur" in
98        -*)
99            COMPREPLY=($(compgen -W "$(_fetch_atest_args)" -- $cur))
100            ;;
101        */*)
102            ;;
103        *)
104            local candidate_args=$(ls; _fetch_testable_modules)
105            COMPREPLY=($(compgen -W "$candidate_args" -- $cur))
106            ;;
107    esac
108
109    case "$prev" in
110        --iterations|--retry-any-failure|--rerun-until-failure)
111            COMPREPLY=(10) ;;
112        --list-modules|-L)
113            # TODO: genetate the list automately when the API is available.
114            COMPREPLY=($(compgen -W "cts vts" -- $cur)) ;;
115        --serial|-s)
116            local adb_devices="$(_fetch_adb_devices)"
117            if [ -n "$adb_devices" ]; then
118                COMPREPLY=($(compgen -W "$(_fetch_adb_devices)" -- $cur))
119            else
120                # Don't complete files/dirs when there'is no devices.
121                compopt -o nospace
122                COMPREPLY=("")
123            fi ;;
124        --test-mapping|-p)
125            local mapping_files="$(_fetch_test_mapping_files)"
126            if [ -n "$mapping_files" ]; then
127                COMPREPLY=($(compgen -W "$mapping_files" -- $cur))
128            else
129                # Don't complete files/dirs when TEST_MAPPING wasn't found.
130                compopt -o nospace
131                COMPREPLY=("")
132            fi ;;
133    esac
134    __ltrim_colon_completions "$cur" "$prev" || true
135    return 0
136}
137
138function _atest_main() {
139    # Only use this in interactive mode.
140    # Warning: below check must be "return", not "exit". "exit" won't break the
141    # build in interactive shell(e.g VM), but will result in build breakage in
142    # non-interactive shell(e.g docker container); therefore, using "return"
143    # adapts both conditions.
144    [[ ! $- =~ 'i' ]] && return 0
145
146    # If required functions are not available, exit gracefully.
147    if ! _atest_completion_ready; then
148        return 0
149    fi
150
151    # Use Py3 as the default interpreter. This script is aiming for being
152    # compatible with both Py2 and Py3.
153    if [ -x "$(which python3)" ]; then
154        PYTHON=$(which python3)
155    elif [ -x "$(which python2)" ]; then
156        PYTHON=$(which python2)
157    else
158        PYTHON="/usr/bin/env python3"
159    fi
160
161    # Complete file/dir name first by using option "nosort".
162    # BASH version <= 4.3 doesn't have nosort option.
163    # Note that nosort has no effect for zsh.
164    local _atest_comp_options="-o default -o nosort"
165    local _atest_executables=(atest atest-dev atest-src atest-py3)
166    for exec in "${_atest_executables[*]}"; do
167        complete -F _atest $_atest_comp_options $exec 2>/dev/null || \
168        complete -F _atest -o default $exec
169    done
170
171    # Install atest-src for the convenience of debugging.
172    local atest_src="$(gettop)/$ATEST_REL_DIR/atest.py"
173    [[ -f "$atest_src" ]] && alias atest-src="$atest_src"
174}
175
176_atest_main
177