1 /*
2  * Copyright (C) 2018 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.util;
18 
19 import com.android.tradefed.build.IBuildInfo;
20 import com.android.tradefed.log.LogUtil.CLog;
21 import com.android.tradefed.targetprep.VtsPythonVirtualenvPreparer;
22 import java.io.ByteArrayOutputStream;
23 import java.io.File;
24 import java.io.IOException;
25 import java.io.OutputStream;
26 
27 /**
28  * A helper class for executing VTS python scripts.
29  */
30 public class VtsPythonRunnerHelper {
31     // The timeout for the runner's teardown prodedure.
32     public static final long TEST_ABORT_TIMEOUT_MSECS = 1000 * 40;
33 
34     static final String PATH = "PATH";
35     static final String PYTHONHOME = "PYTHONHOME";
36     static final String VTS = "vts";
37 
38     // Python virtual environment root path
39     private File mVirtualenvPath;
40     protected IRunUtil mRunUtil;
41 
VtsPythonRunnerHelper(IBuildInfo buildInfo, File workingDir)42     public VtsPythonRunnerHelper(IBuildInfo buildInfo, File workingDir) {
43         this(buildInfo.getBuildAttributes().get(VtsPythonVirtualenvPreparer.VIRTUAL_ENV),
44                 workingDir);
45     }
46 
VtsPythonRunnerHelper(String virtualEnvPath, File workingDir)47     public VtsPythonRunnerHelper(String virtualEnvPath, File workingDir) {
48         this(virtualEnvPath == null ? null : new File(virtualEnvPath), workingDir);
49     }
50 
VtsPythonRunnerHelper(File virtualEnvPath, File workingDir)51     public VtsPythonRunnerHelper(File virtualEnvPath, File workingDir) {
52         this(virtualEnvPath, workingDir, new RunUtil());
53     }
54 
VtsPythonRunnerHelper(File virtualEnvPath, File workingDir, IRunUtil runUtil)55     public VtsPythonRunnerHelper(File virtualEnvPath, File workingDir, IRunUtil runUtil) {
56         mVirtualenvPath = virtualEnvPath;
57         mRunUtil = runUtil;
58         activateVirtualenv(mRunUtil, getPythonVirtualEnv());
59         mRunUtil.setWorkingDir(workingDir);
60     }
61 
62     /**
63      * Run VTS Python runner and handle interrupt from TradeFed.
64      *
65      * @param cmd the command to start VTS Python runner.
66      * @param commandResult the object containing the command result.
67      * @param timeout command timeout value.
68      * @return null if the command terminates or times out; a message string if the command is
69      * interrupted by TradeFed.
70      */
runPythonRunner(String[] cmd, CommandResult commandResult, long timeout)71     public String runPythonRunner(String[] cmd, CommandResult commandResult, long timeout) {
72         OutputStream stdOut = new ByteArrayOutputStream();
73         OutputStream stdErr = new ByteArrayOutputStream();
74         try {
75             return runPythonRunner(cmd, commandResult, timeout, stdOut, stdErr);
76         } finally {
77             try {
78                 stdOut.close();
79                 stdErr.close();
80                 commandResult.setStdout(((ByteArrayOutputStream) stdOut).toString("UTF-8"));
81                 commandResult.setStderr(((ByteArrayOutputStream) stdErr).toString("UTF-8"));
82             } catch (IOException e) {
83                 throw new RuntimeException(e);
84             }
85         }
86     }
87 
88     /**
89      * Run VTS Python runner and handle interrupt from TradeFed.
90      *
91      * @param cmd the command to start VTS Python runner.
92      * @param commandResult the object containing the command result.
93      * @param timeout command timeout value.
94      * @return null if the command terminates or times out; a message string if the command is
95      * interrupted by TradeFed.
96      */
runPythonRunner(String[] cmd, CommandResult commandResult, long timeout, OutputStream stdOut, OutputStream stdErr)97     public String runPythonRunner(String[] cmd, CommandResult commandResult, long timeout,
98             OutputStream stdOut, OutputStream stdErr) {
99         String interruptMessage;
100         CommandStatus commandStatus;
101         CommandResult result;
102         try {
103             result = mRunUtil.runTimedCmd(timeout, stdOut, stdErr, cmd);
104             commandStatus = result.getStatus();
105             interruptMessage = null;
106         } catch (RunInterruptedException e) {
107             CLog.e("Python process is interrupted.");
108             commandStatus = CommandStatus.TIMED_OUT;
109             interruptMessage = (e.getMessage() != null ? e.getMessage() : "");
110         }
111         commandResult.setStatus(commandStatus);
112         return interruptMessage;
113     }
114 
115     /**
116      * Gets python bin directory path.
117      *
118      * This method will check the directory existence.
119      *
120      * @return python bin directory; null if not exist.
121      */
getPythonBinDir(String virtualenvPath)122     public static String getPythonBinDir(String virtualenvPath) {
123         if (virtualenvPath == null) {
124             return null;
125         }
126         String binDirName = EnvUtil.isOnWindows() ? "Scripts" : "bin";
127         File res = new File(virtualenvPath, binDirName);
128         if (!res.exists()) {
129             return null;
130         }
131         return res.getAbsolutePath();
132     }
133 
134     /**
135      * Get python virtualenv path
136      * @return virutalenv path. null if doesn't exist
137      */
getPythonVirtualEnv()138     public String getPythonVirtualEnv() {
139         if (mVirtualenvPath == null) {
140             return null;
141         }
142         return mVirtualenvPath.getAbsolutePath();
143     }
144 
145     /**
146      * Activate virtualenv for a RunUtil.
147      *
148      * This method will check for python bin directory existence
149      *
150      * @param runUtil
151      * @param virtualenvPath
152      */
activateVirtualenv(IRunUtil runUtil, String virtualenvPath)153     public static void activateVirtualenv(IRunUtil runUtil, String virtualenvPath) {
154         String pythonBinDir = getPythonBinDir(virtualenvPath);
155         if (pythonBinDir == null || !new File(pythonBinDir).exists()) {
156             CLog.e("Invalid python virtualenv path. Using python from system path.");
157         } else {
158             String separater = EnvUtil.isOnWindows() ? ";" : ":";
159             runUtil.setEnvVariable(PATH, pythonBinDir + separater + System.getenv().get(PATH));
160             runUtil.setEnvVariable(VtsPythonVirtualenvPreparer.VIRTUAL_ENV, virtualenvPath);
161             runUtil.unsetEnvVariable(PYTHONHOME);
162         }
163     }
164 }
165