1 /*
2  * Copyright (C) 2016 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 
17 package com.android.tradefed.targetprep;
18 
19 import com.android.tradefed.build.IBuildInfo;
20 import com.android.tradefed.config.Option;
21 import com.android.tradefed.config.OptionClass;
22 import com.android.tradefed.device.DeviceNotAvailableException;
23 import com.android.tradefed.device.ITestDevice;
24 import com.android.tradefed.log.LogUtil.CLog;
25 import com.android.tradefed.util.CommandResult;
26 import com.android.tradefed.util.CommandStatus;
27 import com.android.tradefed.util.FileUtil;
28 import com.android.tradefed.util.IRunUtil;
29 import com.android.tradefed.util.RunUtil;
30 
31 import java.io.File;
32 import java.io.IOException;
33 import java.util.ArrayList;
34 import java.util.List;
35 
36 /**
37  * Sets up a Python virtualenv on the host and installs packages. To activate it, the working
38  * directory is changed to the root of the virtualenv.
39  */
40 @OptionClass(alias = "python-venv")
41 public class PythonVirtualenvPreparer extends BaseTargetPreparer {
42 
43     private static final String PIP = "pip";
44     private static final String PATH = "PATH";
45     protected static final String PYTHONPATH = "PYTHONPATH";
46     private static final int BASE_TIMEOUT = 1000 * 60;
47 
48     @Option(name = "venv-dir", description = "path of an existing virtualenv to use")
49     private File mVenvDir = null;
50 
51     @Option(name = "requirements-file", description = "pip-formatted requirements file")
52     private File mRequirementsFile = null;
53 
54     @Option(name = "dep-module", description = "modules which need to be installed by pip")
55     private List<String> mDepModules = new ArrayList<>();
56 
57     IRunUtil mRunUtil = new RunUtil();
58     String mPip = PIP;
59 
60     @Override
setUp(ITestDevice device, IBuildInfo buildInfo)61     public void setUp(ITestDevice device, IBuildInfo buildInfo)
62             throws TargetSetupError, BuildError, DeviceNotAvailableException {
63         if (isDisabled()) {
64             CLog.i("Skipping PythonVirtualenvPreparer");
65             return;
66         }
67         startVirtualenv(buildInfo, device);
68         installDeps(buildInfo, device);
69     }
70 
installDeps(IBuildInfo buildInfo, ITestDevice device)71     protected void installDeps(IBuildInfo buildInfo, ITestDevice device) throws TargetSetupError {
72         boolean hasDependencies = false;
73         if (mRequirementsFile != null) {
74             CommandResult c = mRunUtil.runTimedCmd(BASE_TIMEOUT * 5, mPip,
75                     "install", "-r", mRequirementsFile.getAbsolutePath());
76             if (c.getStatus() != CommandStatus.SUCCESS) {
77                 CLog.e("Installing dependencies from %s failed",
78                         mRequirementsFile.getAbsolutePath());
79                 throw new TargetSetupError("Failed to install dependencies with pip",
80                         device.getDeviceDescriptor());
81             }
82             hasDependencies = true;
83         }
84         if (!mDepModules.isEmpty()) {
85             for (String dep : mDepModules) {
86                 CLog.i("Attempting installation of %s", dep);
87                 CommandResult c = mRunUtil.runTimedCmd(BASE_TIMEOUT * 5, mPip,
88                         "install", dep);
89                 if (c.getStatus() != CommandStatus.SUCCESS) {
90                     CLog.e("Installing %s failed", dep);
91                     throw new TargetSetupError("Failed to install dependencies with pip",
92                             device.getDeviceDescriptor());
93                 }
94                 hasDependencies = true;
95             }
96         }
97         if (!hasDependencies) {
98             CLog.i("No dependencies to install");
99         } else {
100             // make the install directory of new packages available to other classes that
101             // receive the build
102             buildInfo.setFile(PYTHONPATH, new File(mVenvDir,
103                     "local/lib/python2.7/site-packages"),
104                     buildInfo.getBuildId());
105         }
106     }
107 
startVirtualenv(IBuildInfo buildInfo, ITestDevice device)108     protected void startVirtualenv(IBuildInfo buildInfo, ITestDevice device)
109             throws TargetSetupError {
110         if (mVenvDir != null) {
111             CLog.i("Using existing virtualenv based at %s", mVenvDir.getAbsolutePath());
112             activate();
113             return;
114         }
115         try {
116             mVenvDir = FileUtil.createNamedTempDir(buildInfo.getTestTag() + "-virtualenv");
117             mRunUtil.runTimedCmd(BASE_TIMEOUT, "virtualenv", mVenvDir.getAbsolutePath());
118             activate();
119         } catch (IOException e) {
120             CLog.e("Failed to create temp directory for virtualenv");
121             throw new TargetSetupError("Error creating virtualenv", e,
122                     device.getDeviceDescriptor());
123         }
124     }
125 
addDepModule(String module)126     protected void addDepModule(String module) {
127         mDepModules.add(module);
128     }
129 
setRequirementsFile(File f)130     protected void setRequirementsFile(File f) {
131         mRequirementsFile = f;
132     }
133 
activate()134     private void activate() {
135         File binDir = new File(mVenvDir, "bin");
136         mRunUtil.setWorkingDir(binDir);
137         String path = System.getenv(PATH);
138         mRunUtil.setEnvVariable(PATH, binDir + ":" + path);
139         File pipFile = new File(binDir, PIP);
140         pipFile.setExecutable(true);
141         mPip = pipFile.getAbsolutePath();
142     }
143 }