1#!/usr/bin/env python3
2
3# Copyright (C) 2024 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 pathlib
19import subprocess
20import sys
21
22SOURCE_ENVSETUP="source build/make/envsetup.sh && "
23
24def update_display():
25    sys.stderr.write("passed\n")
26
27def go_to_root():
28    while True:
29        if os.path.exists("build/make/envsetup.sh"):
30            return
31        if os.getcwd() == "/":
32            sys.stderr.write("Can't find root of the source tree\n");
33            print("\nFAILED")
34            sys.exit(1)
35        os.chdir("..")
36
37def is_test(name, thing):
38    if not callable(thing):
39        return False
40    if name == "test":
41        return False
42    return name.startswith("test")
43
44
45def test(shell, command, expected_return, expected_stdout, expected_stderr, expected_env):
46    command += "; _rc=$?"
47    for env in expected_env.keys():
48        command += f"; echo ENV: {env}=\\\"${env}\\\""
49    command += "; exit $_rc"
50
51    cmd = [shell, "-c", command]
52    result = subprocess.run(cmd, capture_output=True, text=True)
53
54    status = True
55
56    if result.returncode != expected_return:
57        print()
58        print(f"Expected return code: {expected_return}")
59        print(f"Actual return code:   {result.returncode}")
60        status = False
61
62    printed_stdout = False
63    if expected_stdout and expected_stdout not in result.stdout:
64        print()
65        print(f"Expected stdout to contain:\n{expected_stdout}")
66        print(f"\nActual stdout:\n{result.stdout}")
67        printed_stdout = True
68        status = False
69
70    if expected_stderr and expected_stderr not in result.stderr:
71        print()
72        print(f"Expected stderr to contain:\n{expected_stderr}")
73        print(f"\nActual stderr:\n{result.stderr}")
74        status = False
75
76    env_failure = False
77    for k, v in expected_env.items():
78        if f"{k}=\"{v}\"" not in result.stdout:
79            print()
80            print(f"Expected environment variable {k} to be: {v} --- {k}=\"{v}\"")
81            env_failure = True
82            status = False
83
84    if env_failure and not printed_stdout:
85        print()
86        print("See stdout:")
87        print(result.stdout)
88
89    if not status:
90        print()
91        print("Command to reproduce:")
92        print(command)
93        print()
94
95    return status
96
97NO_LUNCH = {
98    "TARGET_PRODUCT": "",
99    "TARGET_RELEASE": "",
100    "TARGET_BUILD_VARIANT": "",
101}
102
103def test_invalid_lunch_target(shell):
104    return test(shell, SOURCE_ENVSETUP + "lunch invalid-trunk_staging-eng",
105         expected_return=1, expected_stdout=None,
106         expected_stderr="Cannot locate config makefile for product",
107         expected_env=NO_LUNCH)
108
109
110def test_aosp_arm(shell):
111    return test(shell, SOURCE_ENVSETUP + "lunch aosp_arm-trunk_staging-eng",
112         expected_return=0, expected_stdout=None, expected_stderr=None,
113         expected_env={
114            "TARGET_PRODUCT": "aosp_arm",
115            "TARGET_RELEASE": "trunk_staging",
116            "TARGET_BUILD_VARIANT": "eng",
117        })
118
119
120def test_lunch2_empty(shell):
121    return test(shell, SOURCE_ENVSETUP + "lunch2",
122         expected_return=1, expected_stdout=None,
123         expected_stderr="No target specified. See lunch --help",
124         expected_env=NO_LUNCH)
125
126def test_lunch2_four_params(shell):
127    return test(shell, SOURCE_ENVSETUP + "lunch2 a b c d",
128         expected_return=1, expected_stdout=None,
129         expected_stderr="Too many parameters given. See lunch --help",
130         expected_env=NO_LUNCH)
131
132def test_lunch2_aosp_arm(shell):
133    return test(shell, SOURCE_ENVSETUP + "lunch2 aosp_arm",
134         expected_return=0, expected_stdout="=========", expected_stderr=None,
135         expected_env={
136            "TARGET_PRODUCT": "aosp_arm",
137            "TARGET_RELEASE": "trunk_staging",
138            "TARGET_BUILD_VARIANT": "eng",
139        })
140
141def test_lunch2_aosp_arm_trunk_staging(shell):
142    # Somewhat unfortunate because trunk_staging is the only config in
143    # aosp so we can't really test that this isn't just getting the default
144    return test(shell, SOURCE_ENVSETUP + "lunch2 aosp_arm trunk_staging",
145         expected_return=0, expected_stdout="=========", expected_stderr=None,
146         expected_env={
147            "TARGET_PRODUCT": "aosp_arm",
148            "TARGET_RELEASE": "trunk_staging",
149            "TARGET_BUILD_VARIANT": "eng",
150        })
151
152def test_lunch2_aosp_arm_trunk_staging_userdebug(shell):
153    return test(shell, SOURCE_ENVSETUP + "lunch2 aosp_arm trunk_staging userdebug",
154         expected_return=0, expected_stdout="=========", expected_stderr=None,
155         expected_env={
156            "TARGET_PRODUCT": "aosp_arm",
157            "TARGET_RELEASE": "trunk_staging",
158            "TARGET_BUILD_VARIANT": "userdebug",
159        })
160
161def test_list_products(shell):
162    return test(shell, "build/soong/bin/list_products",
163         expected_return=0, expected_stdout="aosp_arm", expected_stderr=None,
164         expected_env=NO_LUNCH)
165
166def test_list_releases_param(shell):
167    return test(shell, "build/soong/bin/list_releases aosp_arm",
168         expected_return=0, expected_stdout="trunk_staging", expected_stderr=None,
169         expected_env=NO_LUNCH)
170
171def test_list_releases_env(shell):
172    return test(shell, "TARGET_PRODUCT=aosp_arm build/soong/bin/list_releases",
173         expected_return=0, expected_stdout="trunk_staging", expected_stderr=None,
174         expected_env=NO_LUNCH)
175
176def test_list_releases_no_product(shell):
177    return test(shell, "build/soong/bin/list_releases",
178         expected_return=1, expected_stdout=None, expected_stderr=None,
179         expected_env=NO_LUNCH)
180
181def test_list_variants(shell):
182    return test(shell, "build/soong/bin/list_variants",
183         expected_return=0, expected_stdout="userdebug", expected_stderr=None,
184         expected_env=NO_LUNCH)
185
186
187def test_get_build_var_in_path(shell):
188    return test(shell, SOURCE_ENVSETUP + "which get_build_var ",
189         expected_return=0, expected_stdout="soong/bin", expected_stderr=None,
190         expected_env=NO_LUNCH)
191
192
193
194TESTS=sorted([(name, thing) for name, thing in locals().items() if is_test(name, thing)])
195
196def main():
197    if any([x.endswith("/soong/bin") for x in os.getenv("PATH").split(":")]):
198        sys.stderr.write("run_envsetup_tests must be run in a shell that has not sourced"
199                + " envsetup.sh\n\nFAILED\n")
200        return 1
201
202    go_to_root()
203
204    tests = TESTS
205    if len(sys.argv) > 1:
206        tests = [(name, func) for name, func in tests if name in sys.argv]
207
208    shells = ["/usr/bin/bash", "/usr/bin/zsh"]
209    total_count = len(tests) * len(shells)
210    index = 1
211    failed_tests = 0
212
213    for name, func in tests:
214        for shell in shells:
215            sys.stdout.write(f"\33[2K\r{index} of {total_count}: {name} in {shell}")
216            passed = func(shell)
217            if not passed:
218                failed_tests += 1
219            index += 1
220
221    if failed_tests > 0:
222        print(f"\n\nFAILED: {failed_tests} of {total_count}")
223        return 1
224    else:
225        print("\n\nSUCCESS")
226        return 0
227
228if __name__ == "__main__":
229    sys.exit(main())
230