xref: /aosp_15_r20/platform_testing/libraries/screenshot/update_goldens2.py (revision dd0948b35e70be4c0246aabd6c72554a5eb8b22a)
1#!/usr/bin/env python3
2
3import argparse
4import os
5import shutil
6
7
8def parse_arguments():
9    """Parses command-line arguments and returns the parsed arguments object."""
10    parser = argparse.ArgumentParser(description="Validate directories and golden files.")
11    parser.add_argument("--source-directory", required=True, help="Path to the source directory.")
12    parser.add_argument("--android-build-top", required=True,
13                        help="Path to the Android build directory.")
14    return parser.parse_args()
15
16
17def validate_directories(args):
18    """Validates the provided source and Android directory arguments and returns their paths if valid."""
19    if not args.source_directory or not args.android_build_top:
20        print("Error: Both --source-directory and --android-build-top arguments are required.")
21        return None
22
23    source_dir = args.source_directory
24    android_dir = args.android_build_top
25
26    is_source_dir_valid = os.path.isdir(source_dir)
27    is_android_dir_valid = os.path.isdir(android_dir)
28
29    if not is_source_dir_valid:
30        print(f"Error: Source directory does not exist: {source_dir}")
31
32    if not is_android_dir_valid:
33        print(f"Error: Android build directory does not exist: {android_dir}")
34
35    if not is_source_dir_valid or not is_android_dir_valid:
36        return None
37
38    return [source_dir, android_dir]
39
40
41def find_golden_files(source_dir):
42    """Finds golden files within the source directory and returns their filenames or handles errors."""
43    golden_files = []
44    for root, _, files in os.walk(source_dir):
45        for file in files:
46            if "_goldResult_" in file and file.endswith(".textproto"):
47                golden_files.append(file)
48
49    if not golden_files:
50        print("Error: No golden files found in the source directory.")
51        return None
52
53    return golden_files
54
55
56def validate_protos(source_dir, proto_files):
57    """Validates proto files, extracts image locations, and handles errors."""
58    all_image_locations = []
59    for filename in proto_files:
60        required_image_locations = {"image_location_diff": False, "image_location_golden": False,
61                                    "image_location_reference": False, "image_location_test": False}
62        found_image_locations = {}
63
64        with open(os.path.join(source_dir, filename), 'r') as file:
65            for line in file:
66                for location in required_image_locations:
67                    if line.startswith(location):
68                        if required_image_locations[location]:
69                            print(f"Error: Duplicate '{location}' entry found in {filename}.")
70                            return None
71                        required_image_locations[location] = True
72                        found_image_locations[location] = line.split(": ")[1].strip().strip('"')
73                        break
74
75        if not all(required_image_locations.values()):
76            missing_keys = [key for key, found in required_image_locations.items() if not found]
77            print(f"Error: Missing image location(s) in {filename}: {', '.join(missing_keys)}")
78            return None
79        else:
80            print(f"Proto file {filename} is valid.")
81            all_image_locations.append(found_image_locations)
82    return all_image_locations
83
84
85def validate_test_images(source_dir, all_image_locations):
86    """Validates if PNG files exist in the source directory based on provided image locations."""
87    for image_locations in all_image_locations:
88        for location, path in image_locations.items():
89            if location != "image_location_golden":
90                base_name = os.path.splitext(os.path.basename(path))[0]
91                for root, _, files in os.walk(source_dir):
92                    for file in files:
93                        if file.startswith(base_name) and file.endswith(".png"):
94                            image_locations[location] = os.path.join(root, file)
95                            break
96                    else:
97                        print(f"Error: No PNG file found matching {path} in {source_dir}")
98                        return None
99        filename_without_ext = \
100        os.path.splitext(os.path.basename(image_locations["image_location_golden"]))[0]
101        print(f"Golden {filename_without_ext} is valid.")
102    return all_image_locations
103
104
105def update_goldens(android_dir, updated_image_locations_list):
106    """Copies updated 'image_location_test' images to their 'image_location_golden' paths in the android_build_top directory."""
107    for image_locations in updated_image_locations_list:
108        test_image_path = image_locations["image_location_test"]
109        golden_image_path = os.path.join(android_dir, image_locations["image_location_golden"])
110
111        try:
112            shutil.copy2(test_image_path, golden_image_path)
113            print(f"Updated golden image: {golden_image_path}")
114        except IOError as e:
115            print(f"Error updating golden image: {e}")
116
117
118def main():
119    args = parse_arguments()
120
121    directories = validate_directories(args)
122    if directories is None:
123        return
124
125    source_dir, android_dir = directories
126
127    proto_files = find_golden_files(source_dir)
128    if proto_files is None:
129        return
130
131    all_image_locations = validate_protos(source_dir, proto_files)
132    if all_image_locations is None:
133        return
134
135    updated_image_locations_list = validate_test_images(source_dir, all_image_locations)
136    if updated_image_locations_list is None:
137        return
138
139    update_goldens(android_dir, updated_image_locations_list)
140
141
142if __name__ == "__main__":
143    main()
144