harnesses/tradefed/src/com/android/tradefed/targetprep/VtsPythonVirtualenvPreparer.java

/*
 * Copyright (C) 2016 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.android.tradefed.targetprep;

import com.android.tradefed.build.IBuildInfo;
import com.android.tradefed.config.Option;
import com.android.tradefed.config.OptionClass;
import com.android.tradefed.device.DeviceNotAvailableException;
import com.android.tradefed.device.ITestDevice;
import com.android.tradefed.log.LogUtil.CLog;
import com.android.tradefed.util.CommandResult;
import com.android.tradefed.util.CommandStatus;
import com.android.tradefed.util.FileUtil;
import com.android.tradefed.util.IRunUtil;
import com.android.tradefed.util.RunUtil;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.TreeSet;
import java.util.Collection;

/**
 * Sets up a Python virtualenv on the host and installs packages. To activate it, the working
 * directory is changed to the root of the virtualenv.
 *
 * This's a fork of PythonVirtualenvPreparer and is forked in order to simplify the change
 * deployment process and reduce the deployment time, which are critical for VTS services.
 * That means changes here will be upstreamed gradually.
 */
@OptionClass(alias = "python-venv")
public class VtsPythonVirtualenvPreparer implements ITargetPreparer, ITargetCleaner {

    private static final String PIP = "pip";
    private static final String PATH = "PATH";
    private static final String OS_NAME = "os.name";
    private static final String WINDOWS = "Windows";
    protected static final String PYTHONPATH = "PYTHONPATH";
    protected static final String VIRTUAL_ENV_PATH = "VIRTUALENVPATH";
    private static final int BASE_TIMEOUT = 1000 * 60;
    private static final String[] DEFAULT_DEP_MODULES = {
            "future", "futures", "enum", "protobuf", "requests", "httplib2",
            "google-api-python-client", "oauth2client"};

    @Option(name = "venv-dir", description = "path of an existing virtualenv to use")
    private File mVenvDir = null;

    @Option(name = "requirements-file", description = "pip-formatted requirements file")
    private File mRequirementsFile = null;

    @Option(name = "script-file", description = "scripts which need to be executed in advance")
    private Collection<String> mScriptFiles = new TreeSet<>();

    @Option(name = "dep-module", description = "modules which need to be installed by pip")
    private Collection<String> mDepModules = new TreeSet<>(Arrays.asList(DEFAULT_DEP_MODULES));

    IRunUtil mRunUtil = new RunUtil();
    String mPip = PIP;

    /**
     * {@inheritDoc}
     */
    @Override
    public void setUp(ITestDevice device, IBuildInfo buildInfo)
            throws TargetSetupError, BuildError, DeviceNotAvailableException {
        startVirtualenv(buildInfo);
        installDeps(buildInfo);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void tearDown(ITestDevice device, IBuildInfo buildInfo, Throwable e)
            throws DeviceNotAvailableException {
        if (mVenvDir != null) {
            FileUtil.recursiveDelete(mVenvDir);
            CLog.i("Deleted the virtual env's temp working dir, %s.", mVenvDir);
            mVenvDir = null;
        }
    }

    protected void installDeps(IBuildInfo buildInfo) throws TargetSetupError {
        boolean hasDependencies = false;
        if (!mScriptFiles.isEmpty()) {
            for (String scriptFile : mScriptFiles) {
                CLog.i("Attempting to execute a script, %s", scriptFile);
                CommandResult c = mRunUtil.runTimedCmd(BASE_TIMEOUT * 5, scriptFile);
                if (c.getStatus() != CommandStatus.SUCCESS) {
                    CLog.e("Executing script %s failed", scriptFile);
                    throw new TargetSetupError("Failed to source a script");
                }
            }
        }
        if (mRequirementsFile != null) {
            CommandResult c = mRunUtil.runTimedCmd(BASE_TIMEOUT * 5, mPip,
                    "install", "-r", mRequirementsFile.getAbsolutePath());
            if (c.getStatus() != CommandStatus.SUCCESS) {
                CLog.e("Installing dependencies from %s failed",
                        mRequirementsFile.getAbsolutePath());
                throw new TargetSetupError("Failed to install dependencies with pip");
            }
            hasDependencies = true;
        }
        if (!mDepModules.isEmpty()) {
            for (String dep : mDepModules) {
                CLog.i("Attempting installation of %s", dep);
                CommandResult result = mRunUtil.runTimedCmd(
                        BASE_TIMEOUT * 5, mPip, "install", dep);
                CLog.i(String.format(
                    "Result %s. stdout: %s, stderr: %s",
                    result.getStatus(), result.getStdout(), result.getStderr()));
                if (result.getStatus() != CommandStatus.SUCCESS) {
                    CLog.e("Installing %s failed.", dep);
                    CLog.i("Attempting to upgrade %s", dep);
                    result = mRunUtil.runTimedCmd(
                            BASE_TIMEOUT * 5, mPip, "install", "--upgrade", dep);
                    if (result.getStatus() != CommandStatus.SUCCESS) {
                        throw new TargetSetupError(String.format(
                            "Failed to install dependencies with pip. "
                                    + "Result %s. stdout: %s, stderr: %s",
                            result.getStatus(), result.getStdout(), result.getStderr()));
                    } else {
                        CLog.i(String.format(
                            "Result %s. stdout: %s, stderr: %s",
                            result.getStatus(), result.getStdout(), result.getStderr()));
                    }
                }
                hasDependencies = true;
            }
        }
        if (!hasDependencies) {
            CLog.i("No dependencies to install");
        } else {
            // make the install directory of new packages available to other classes that
            // receive the build
            buildInfo.setFile(PYTHONPATH, new File(mVenvDir,
                    "local/lib/python2.7/site-packages"),
                    buildInfo.getBuildId());
        }
    }

    protected void startVirtualenv(IBuildInfo buildInfo) throws TargetSetupError {
        if (mVenvDir != null) {
            CLog.i("Using existing virtualenv based at %s", mVenvDir.getAbsolutePath());
            activate();
            return;
        }
        try {
            mVenvDir = buildInfo.getFile(VIRTUAL_ENV_PATH);
            if (mVenvDir == null) {
                mVenvDir = FileUtil.createTempDir(buildInfo.getTestTag() + "-virtualenv");
            }
            String virtualEnvPath = mVenvDir.getAbsolutePath();
            mRunUtil.runTimedCmd(BASE_TIMEOUT, "virtualenv", virtualEnvPath);
            CLog.i(VIRTUAL_ENV_PATH + " = " + virtualEnvPath + "\n");
            buildInfo.setFile(VIRTUAL_ENV_PATH, new File(virtualEnvPath),
                              buildInfo.getBuildId());
            activate();
        } catch (IOException e) {
            CLog.e("Failed to create temp directory for virtualenv");
            throw new TargetSetupError("Error creating virtualenv", e);
        }
    }

    protected void addDepModule(String module) {
        mDepModules.add(module);
    }

    protected void setRequirementsFile(File f) {
        mRequirementsFile = f;
    }

    /**
     * This method returns whether the OS is Windows.
     */
    private static boolean isOnWindows() {
        return System.getProperty(OS_NAME).contains(WINDOWS);
    }

    private void activate() {
        File binDir = new File(mVenvDir, isOnWindows() ? "Scripts" : "bin");
        mRunUtil.setWorkingDir(binDir);
        String path = System.getenv(PATH);
        mRunUtil.setEnvVariable(PATH, binDir + File.pathSeparator + path);
        File pipFile = new File(binDir, PIP);
        pipFile.setExecutable(true);
        mPip = pipFile.getAbsolutePath();
    }
}