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/tradefederation/core/atest"
16
17_fetch_testable_modules() {
18    [[ -z $ANDROID_BUILD_TOP ]] && return 0
19    export ATEST_DIR="$ANDROID_BUILD_TOP/$ATEST_REL_DIR"
20    $PYTHON - << END
21import os
22import pickle
23import sys
24
25sys.path.append(os.getenv('ATEST_DIR'))
26import constants
27
28if os.path.isfile(constants.MODULE_INDEX):
29    with open(constants.MODULE_INDEX, 'rb') as cache:
30        try:
31            print("\n".join(pickle.load(cache, encoding="utf-8")))
32        except:
33            print("\n".join(pickle.load(cache)))
34else:
35    print("")
36END
37    unset ATEST_DIR
38}
39
40# This function invoke get_args() and return each item
41# of the list for tab completion candidates.
42_fetch_atest_args() {
43    [[ -z $ANDROID_BUILD_TOP ]] && return 0
44    export ATEST_DIR="$ANDROID_BUILD_TOP/$ATEST_REL_DIR"
45    $PYTHON - << END
46import os
47import sys
48
49atest_dir = os.path.join(os.getenv('ATEST_DIR'))
50sys.path.append(atest_dir)
51
52import atest_arg_parser
53
54parser = atest_arg_parser.AtestArgParser()
55parser.add_atest_args()
56print("\n".join(parser.get_args()))
57END
58    unset ATEST_DIR
59}
60
61# This function returns devices recognised by adb.
62_fetch_adb_devices() {
63    while read dev; do echo $dev | awk '{print $1}'; done < <(adb devices | egrep -v "^List|^$"||true)
64}
65
66# This function returns all paths contain TEST_MAPPING.
67_fetch_test_mapping_files() {
68    [[ -z $ANDROID_BUILD_TOP ]] && return 0
69    find -maxdepth 5 -type f -name TEST_MAPPING |sed 's/^.\///g'| xargs dirname 2>/dev/null
70}
71
72# The main tab completion function.
73_atest() {
74    # Not support completion on Darwin since the bash version of it
75    # is too old to fully support useful built-in commands/functions
76    # such as compopt, _get_comp_words_by_ref and __ltrim_colon_completions.
77    [[ "$(uname -s)" == "Darwin" ]] && return 0
78
79    local cur prev
80    COMPREPLY=()
81    cur="${COMP_WORDS[COMP_CWORD]}"
82    prev="${COMP_WORDS[COMP_CWORD-1]}"
83    _get_comp_words_by_ref -n : cur prev || true
84
85    case "$cur" in
86        -*)
87            COMPREPLY=($(compgen -W "$(_fetch_atest_args)" -- $cur))
88            ;;
89        */*)
90            ;;
91        *)
92            local candidate_args=$(ls; _fetch_testable_modules)
93            COMPREPLY=($(compgen -W "$candidate_args" -- $cur))
94            ;;
95    esac
96
97    case "$prev" in
98        --iterations|--rerun-until-failure|--retry-any-failure)
99            COMPREPLY=(10) ;;
100        --list-modules|-L)
101            # TODO: genetate the list automately when the API is available.
102            COMPREPLY=($(compgen -W "cts vts" -- $cur)) ;;
103        --serial|-s)
104            local adb_devices="$(_fetch_adb_devices)"
105            if [ -n "$adb_devices" ]; then
106                COMPREPLY=($(compgen -W "$(_fetch_adb_devices)" -- $cur))
107            else
108                # Don't complete files/dirs when there'is no devices.
109                compopt -o nospace
110                COMPREPLY=("")
111            fi ;;
112        --test-mapping|-p)
113            local mapping_files="$(_fetch_test_mapping_files)"
114            if [ -n "$mapping_files" ]; then
115                COMPREPLY=($(compgen -W "$mapping_files" -- $cur))
116            else
117                # Don't complete files/dirs when TEST_MAPPING wasn't found.
118                compopt -o nospace
119                COMPREPLY=("")
120            fi ;;
121    esac
122    __ltrim_colon_completions "$cur" "$prev" || true
123    return 0
124}
125
126function _atest_main() {
127    # Only use this in interactive mode.
128    # Warning: below check must be "return", not "exit". "exit" won't break the
129    # build in interactive shell(e.g VM), but will result in build breakage in
130    # non-interactive shell(e.g docker container); therefore, using "return"
131    # adapts both conditions.
132    [[ ! $- =~ 'i' ]] && return 0
133
134    # Use Py2 as the default interpreter. This script is aiming for being
135    # compatible with both Py2 and Py3.
136    if [ -x "$(which python)" ]; then
137        PYTHON=$(which python)
138    elif [ -x "$(which python3)" ]; then
139        PYTHON=$(which python3)
140    else
141        PYTHON="/usr/bin/env python"
142    fi
143
144    # Complete file/dir name first by using option "nosort".
145    # BASH version <= 4.3 doesn't have nosort option.
146    # Note that nosort has no effect for zsh.
147    local _atest_comp_options="-o default -o nosort"
148    local _atest_executables=(atest atest-dev atest-src)
149    for exec in "${_atest_executables[*]}"; do
150        complete -F _atest $_atest_comp_options $exec 2>/dev/null || \
151        complete -F _atest -o default $exec
152    done
153
154    # Install atest-src for the convenience of debugging.
155    local atest_src="$(gettop)/$ATEST_REL_DIR/atest.py"
156    [[ -f "$atest_src" ]] && alias atest-src="$atest_src"
157}
158
159_atest_main
160