xref: /aosp_15_r20/platform_testing/libraries/motion/golden_updater/watch_local_tests.py (revision dd0948b35e70be4c0246aabd6c72554a5eb8b22a)
1#!/usr/bin/env python3
2
3#
4# Copyright 2024, The Android Open Source Project
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10#     http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17#
18
19import http.server
20import socketserver
21import json
22import re
23import urllib.parse
24from os import path
25import socket
26import argparse
27import os
28import subprocess
29import sys
30import tempfile
31import webbrowser
32import mimetypes
33import hashlib
34import shutil
35import secrets
36import datetime
37
38
39from collections import defaultdict
40
41
42def main():
43    parser = argparse.ArgumentParser(
44        "Watches a connected device for golden file updates."
45    )
46
47    parser.add_argument(
48        "--port",
49        default=find_free_port(),
50        type=int,
51        help="Port to run test at watcher web UI on.",
52    )
53    parser.add_argument(
54        "--serial",
55        default=os.environ.get("ANDROID_SERIAL"),
56        help="The ADB device serial to pull goldens from.",
57    )
58
59    parser.add_argument(
60        "--android_build_top",
61        default=os.environ.get("ANDROID_BUILD_TOP"),
62        help="The root directory of the android checkout.",
63    )
64
65    parser.add_argument(
66        "--client_url",
67        default="http://motion.teams.x20web.corp.google.com/",
68        help="The URL where the client app is deployed.",
69    )
70
71    args = parser.parse_args()
72
73    if args.android_build_top is None or not os.path.exists(args.android_build_top):
74        print("ANDROID_BUILD_TOP not set. Have you sourced envsetup.sh?")
75        sys.exit(1)
76
77    serial = args.serial
78    if not serial:
79        devices_response = subprocess.run(
80            ["adb", "devices"], check=True, capture_output=True
81        ).stdout.decode("utf-8")
82        lines = [s for s in devices_response.splitlines() if s.strip()]
83
84        if len(lines) == 1:
85            print("no adb devices found")
86            sys.exit(1)
87
88        if len(lines) > 2:
89            print("multiple adb devices found, specify --serial")
90            sys.exit(1)
91
92        serial = lines[1].split("\t")[0]
93
94    adb_client = AdbClient(serial)
95    if not adb_client.run_as_root():
96        sys.exit(1)
97
98    global android_build_top
99    android_build_top = args.android_build_top
100
101    with tempfile.TemporaryDirectory() as tmpdir:
102        global golden_watcher, this_server_address
103        golden_watcher = GoldenFileWatcher(tmpdir, adb_client)
104
105        this_server_address = f"http://localhost:{args.port}"
106
107        with socketserver.TCPServer(
108            ("localhost", args.port), WatchWebAppRequestHandler, golden_watcher
109        ) as httpd:
110            uiAddress = f"{args.client_url}?token={secret_token}&port={args.port}"
111            print(f"Open UI at {uiAddress}")
112            webbrowser.open(uiAddress)
113            try:
114                httpd.serve_forever()
115            except KeyboardInterrupt:
116                httpd.shutdown()
117                print("Shutting down")
118
119
120GOLDEN_ACCESS_TOKEN_HEADER = "Golden-Access-Token"
121GOLDEN_ACCESS_TOKEN_LOCATION = os.path.expanduser("~/.config/motion-golden/.token")
122
123secret_token = None
124android_build_top = None
125golden_watcher = None
126this_server_address = None
127
128
129class WatchWebAppRequestHandler(http.server.BaseHTTPRequestHandler):
130
131    def __init__(self, *args, **kwargs):
132        self.root_directory = path.abspath(path.dirname(__file__))
133        super().__init__(*args, **kwargs)
134
135    def verify_access_token(self):
136        token = self.headers.get(GOLDEN_ACCESS_TOKEN_HEADER)
137        if not token or token != secret_token:
138            self.send_response(403, "Bad authorization token!")
139            return False
140
141        return True
142
143    def do_OPTIONS(self):
144        self.send_response(200)
145        self.send_header("Allow", "GET,POST,PUT")
146        self.add_standard_headers()
147        self.end_headers()
148        self.wfile.write(b"GET,POST,PUT")
149
150    def do_GET(self):
151
152        parsed = urllib.parse.urlparse(self.path)
153
154        if parsed.path == "/service/list":
155            self.service_list_goldens()
156            return
157        elif parsed.path.startswith("/golden/"):
158            requested_file_start_index = parsed.path.find("/", len("/golden/") + 1)
159            requested_file = parsed.path[requested_file_start_index + 1 :]
160            print(requested_file)
161            self.serve_file(golden_watcher.temp_dir, requested_file)
162            return
163        elif parsed.path.startswith("/expected/"):
164            golden_id = parsed.path[len("/expected/") :]
165            print(golden_id)
166
167            goldens = golden_watcher.cached_goldens.values()
168            for golden in goldens:
169                if golden.id != golden_id:
170                    continue
171
172                self.serve_file(
173                    android_build_top, golden.golden_repo_path, "application/json"
174                )
175                return
176
177        self.send_error(404)
178
179    def do_POST(self):
180        if not self.verify_access_token():
181            return
182
183        content_type = self.headers.get("Content-Type")
184
185        # refuse to receive non-json content
186        if content_type != "application/json":
187            self.send_response(400)
188            return
189
190        length = int(self.headers.get("Content-Length"))
191        message = json.loads(self.rfile.read(length))
192
193        parsed = urllib.parse.urlparse(self.path)
194        if parsed.path == "/service/refresh":
195            self.service_refresh_goldens(message["clear"])
196        else:
197            self.send_error(404)
198
199    def do_PUT(self):
200        if not self.verify_access_token():
201            return
202
203        parsed = urllib.parse.urlparse(self.path)
204        query_params = urllib.parse.parse_qs(parsed.query)
205
206        if parsed.path == "/service/update":
207            self.service_update_golden(query_params["id"][0])
208        else:
209            self.send_error(404)
210
211    def serve_file(self, root_directory, file_relative_to_root, mime_type=None):
212        resolved_path = path.abspath(path.join(root_directory, file_relative_to_root))
213
214        print(resolved_path)
215        print(root_directory)
216
217        if path.commonprefix(
218            [resolved_path, root_directory]
219        ) == root_directory and path.isfile(resolved_path):
220            self.send_response(200)
221            self.send_header(
222                "Content-type", mime_type or mimetypes.guess_type(resolved_path)[0]
223            )
224            self.add_standard_headers()
225            self.end_headers()
226            with open(resolved_path, "rb") as f:
227                self.wfile.write(f.read())
228
229        else:
230            self.send_error(404)
231
232    def service_list_goldens(self):
233        if not self.verify_access_token():
234            return
235
236        goldens_list = []
237
238        for golden in golden_watcher.cached_goldens.values():
239
240            golden_data = {}
241            golden_data["id"] = golden.id
242            golden_data["result"] = golden.result
243            golden_data["label"] = golden.golden_identifier
244            golden_data["goldenRepoPath"] = golden.golden_repo_path
245            golden_data["updated"] = golden.updated
246            golden_data["testClassName"] = golden.test_class_name
247            golden_data["testMethodName"] = golden.test_method_name
248            golden_data["testTime"] = golden.test_time
249
250            golden_data["actualUrl"] = (
251                f"{this_server_address}/golden/{golden.checksum}/{golden.local_file[len(golden_watcher.temp_dir) + 1 :]}"
252            )
253            expected_file = path.join(android_build_top, golden.golden_repo_path)
254            if os.path.exists(expected_file):
255                golden_data["expectedUrl"] = (
256                    f"{this_server_address}/expected/{golden.id}"
257                )
258
259            golden_data["videoUrl"] = (
260                f"{this_server_address}/golden/{golden.checksum}/{golden.video_location}"
261            )
262
263            goldens_list.append(golden_data)
264
265        self.send_json(goldens_list)
266
267    def service_refresh_goldens(self, clear):
268        if clear:
269            golden_watcher.clean()
270        golden_watcher.refresh_golden_files()
271        self.service_list_goldens()
272
273    def service_update_golden(self, id):
274        goldens = golden_watcher.cached_goldens.values()
275        for golden in goldens:
276            if golden.id != id:
277                print("skip", golden.id)
278                continue
279
280            shutil.copyfile(
281                golden.local_file,
282                path.join(android_build_top, golden.golden_repo_path),
283            )
284
285            golden.updated = True
286            self.send_json({"result": "OK"})
287            return
288
289        self.send_error(400)
290
291    def send_json(self, data):
292        # Replace this with code that generates your JSON data
293        data_encoded = json.dumps(data).encode("utf-8")
294        self.send_response(200)
295        self.send_header("Content-type", "application/json")
296        self.add_standard_headers()
297        self.end_headers()
298        self.wfile.write(data_encoded)
299
300    def add_standard_headers(self):
301        self.send_header("Access-Control-Allow-Origin", "*")
302        self.send_header("Access-Control-Allow-Methods", "POST, PUT, GET, OPTIONS")
303        self.send_header(
304            "Access-Control-Allow-Headers",
305            GOLDEN_ACCESS_TOKEN_HEADER
306            + ", Content-Type, Content-Length, Range, Accept-ranges",
307        )
308        # Accept-ranges: bytes is needed for chrome to allow seeking the
309        # video. At this time, won't handle ranges on subsequent gets,
310        # but that is likely OK given the size of these videos and that
311        # its local only.
312        self.send_header("Accept-ranges", "bytes")
313
314
315class GoldenFileWatcher:
316
317    def __init__(self, temp_dir, adb_client):
318        self.temp_dir = temp_dir
319        self.adb_client = adb_client
320
321        # name -> CachedGolden
322        self.cached_goldens = {}
323        self.refresh_golden_files()
324
325    def clean(self):
326        self.cached_goldens = {}
327
328    def refresh_golden_files(self):
329        command = f"find /data/user/0/ -type f -name *.actual.json"
330        updated_goldens = self.run_adb_command(["shell", command]).splitlines()
331        print(f"Updating goldens - found {len(updated_goldens)} files")
332
333        for golden_remote_file in updated_goldens:
334            local_file = self.adb_pull(golden_remote_file)
335
336            golden = CachedGolden(golden_remote_file, local_file)
337            if golden.video_location:
338                self.adb_pull_image(golden.device_local_path, golden.video_location)
339
340            self.cached_goldens[golden_remote_file] = golden
341
342    def adb_pull(self, remote_file):
343        local_file = os.path.join(self.temp_dir, os.path.basename(remote_file))
344        self.run_adb_command(["pull", remote_file, local_file])
345        self.run_adb_command(["shell", "rm", remote_file])
346        return local_file
347
348    def adb_pull_image(self, remote_dir, remote_file):
349        remote_path = os.path.join(remote_dir, remote_file)
350        local_path = os.path.join(self.temp_dir, remote_file)
351        os.makedirs(os.path.dirname(local_path), exist_ok=True)
352        self.run_adb_command(["pull", remote_path, local_path])
353        self.run_adb_command(["shell", "rm", remote_path])
354        return local_path
355
356    def run_adb_command(self, args):
357        return self.adb_client.run_adb_command(args)
358
359
360class CachedGolden:
361
362    def __init__(self, remote_file, local_file):
363        self.id = hashlib.md5(remote_file.encode("utf-8")).hexdigest()
364        self.remote_file = remote_file
365        self.local_file = local_file
366        self.updated = False
367        self.test_time = datetime.datetime.now().isoformat()
368        # Checksum is the time the test data was loaded, forcing unique URLs
369        # every time the golden is reloaded
370        self.checksum = hashlib.md5(self.test_time.encode("utf-8")).hexdigest()
371
372        motion_golden_data = None
373        with open(local_file, "r") as json_file:
374            motion_golden_data = json.load(json_file)
375        metadata = motion_golden_data["//metadata"]
376
377        self.result = metadata["result"]
378        self.golden_repo_path = metadata["goldenRepoPath"]
379        self.golden_identifier = metadata["goldenIdentifier"]
380        self.test_class_name = metadata["testClassName"]
381        self.test_method_name = metadata["testMethodName"]
382        self.device_local_path = metadata["deviceLocalPath"]
383        self.video_location = None
384        if "videoLocation" in metadata:
385            self.video_location = metadata["videoLocation"]
386
387        with open(local_file, "w") as json_file:
388            del motion_golden_data["//metadata"]
389            json.dump(motion_golden_data, json_file, indent=2)
390
391
392class AdbClient:
393    def __init__(self, adb_serial):
394        self.adb_serial = adb_serial
395
396    def run_as_root(self):
397        root_result = self.run_adb_command(["root"])
398        if "restarting adbd as root" in root_result:
399            self.wait_for_device()
400            return True
401        if "adbd is already running as root" in root_result:
402            return True
403
404        print(f"run_as_root returned [{root_result}]")
405
406        return False
407
408    def wait_for_device(self):
409        self.run_adb_command(["wait-for-device"])
410
411    def run_adb_command(self, args):
412        command = ["adb"]
413        command += ["-s", self.adb_serial]
414        command += args
415        return subprocess.run(command, check=True, capture_output=True).stdout.decode(
416            "utf-8"
417        )
418
419
420def find_free_port():
421    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
422        s.bind(("", 0))  # Bind to a random free port provided by the OS
423        return s.getsockname()[1]  # Get the port number
424
425
426def get_token() -> str:
427    try:
428        with open(GOLDEN_ACCESS_TOKEN_LOCATION, "r") as token_file:
429            token = token_file.readline()
430            return token
431    except IOError:
432        token = secrets.token_hex(32)
433        os.makedirs(os.path.dirname(GOLDEN_ACCESS_TOKEN_LOCATION), exist_ok=True)
434        try:
435            with open(GOLDEN_ACCESS_TOKEN_LOCATION, "w") as token_file:
436                token_file.write(token)
437            os.chmod(GOLDEN_ACCESS_TOKEN_LOCATION, 0o600)
438        except IOError:
439            print(
440                "Unable to save persistent token {} to {}".format(
441                    token, GOLDEN_ACCESS_TOKEN_LOCATION
442                )
443            )
444        return token
445
446
447if __name__ == "__main__":
448    secret_token = get_token()
449    main()
450