1#
2# Copyright (C) 2017 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import logging
18import os
19
20from host_controller.build import build_provider
21from vts.utils.python.build.api import artifact_fetcher
22
23
24class BuildProviderAB(build_provider.BuildProvider):
25    """A build provider for Android Build (AB)."""
26
27    def __init__(self):
28        super(BuildProviderAB, self).__init__()
29        if 'run_ab_key' in os.environ:
30            logging.info(
31                "For AB, use the key at %s", os.environ['run_ab_key'])
32            self._artifact_fetcher = artifact_fetcher.AndroidBuildClient(
33                os.environ['run_ab_key'])
34        else:
35            self._artifact_fetcher = None
36
37    def GetLatestBuildId(self, branch, target):
38        """Get the latest build id.
39
40        Args:
41            branch: string, android branch to pull resource from.
42            target: string, build target name.
43
44        Returns:
45            string, latest build id. None if _artifact_fetcher is not initialized.
46        """
47        if not self._artifact_fetcher:
48            return None
49
50        recent_build_ids = self._artifact_fetcher.ListBuildIds(
51            branch, target)
52
53        return recent_build_ids[0]
54
55    def Fetch(self,
56              branch,
57              target,
58              artifact_name,
59              build_id="latest",
60              full_device_images=False):
61        """Fetches Android device artifact file(s) from Android Build.
62
63        Args:
64            branch: string, android branch to pull resource from.
65            target: string, build target name.
66            artifact_name: string, file name.
67            build_id: string, ID of the build or latest.
68
69        Returns:
70            a dict containing the device image info.
71            a dict containing the test suite package info.
72            a dict containing the artifact info.
73        """
74        fetch_info = {}
75        fetch_info["build_id"] = None
76
77        if not self._artifact_fetcher:
78            return self.GetDeviceImage(), self.GetTestSuitePackage(), fetch_info
79
80        if build_id == "latest":
81            build_id = self.GetLatestBuildId(branch, target)
82        fetch_info["build_id"] = build_id
83
84        if "{build_id}" in artifact_name:
85            artifact_name = artifact_name.replace("{build_id}", build_id)
86
87        dest_filepath = os.path.join(self.tmp_dirpath, artifact_name)
88        self._artifact_fetcher.DownloadArtifactToFile(
89            branch, target, build_id, artifact_name,
90            dest_filepath=dest_filepath)
91
92        self.SetFetchedFile(dest_filepath, full_device_images)
93
94        return self.GetDeviceImage(), self.GetTestSuitePackage(), fetch_info
95