vpython: Add VirtualEnv creation package.

vpython/venv/venv_resources_test.go

+// Copyright 2017 The LUCI Authors. All rights reserved.
+// Use of this source code is governed under the Apache License, Version 2.0
+// that can be found in the LICENSE file.
+
+package venv
+
+import (
+        "archive/zip"
+        "crypto/sha256"
+        "encoding/hex"
+        "hash"
+        "io"
+        "net/http"
+        "os"
+        "path/filepath"
+        "testing"
+        "time"
+
+        "github.com/danjacques/gofslock/fslock"
+        "golang.org/x/net/context"
+
+        "github.com/luci/luci-go/cipd/client/cipd"
+        "github.com/luci/luci-go/common/errors"
+        "github.com/luci/luci-go/common/system/filesystem"
+        "github.com/luci/luci-go/common/testing/testfs"
+        "github.com/luci/luci-go/hardcoded/chromeinfra"
+        "github.com/luci/luci-go/vpython/api/vpython"
+        "github.com/luci/luci-go/vpython/python"
+        "github.com/luci/luci-go/vpython/wheel"
+)
+
+const testDataDir = "test_data"
+
+// remoteFiles is the set of remote files to acquire.
+var remoteFiles = []struct {
+        // install installs this remote file into the test environment.
+        install func(te *testingLoader, path string)
+
+        // name is the name of the file.
+        name string
+        // contentHash is the SHA256 has of the content.
+        contentHash string
+
+        // cipdPackage, if not empty, is the name of the CIPD package that contains
+        // this file.
+        cipdPackage string
+        // cipdVersion is the version string of the CIPD package.
+        cipdVersion string
+
+        // urls, if not empty, is a set of remote URLs where this file can be
+        // downloaded from.
+        urls []string
+}{
+        {
+                install:     func(tl *testingLoader, path string) { tl.virtualEnvZIP = path },
+                name:        "virtualenv-15.1.0.zip",
+                contentHash: "f7682a57c98a10d32474b4c1df75478dea9a0802c140335c0269a6ec3af46201",
+                cipdPackage: "infra/test-data/vpython/virtualenv",
+                cipdVersion: "version:15.1.0",
+                urls: []string{
+                        "https://github.com/pypa/virtualenv/archive/15.1.0.zip",
+                },
+        },
+}
+
+// testingLoader is a map of a CIPD package name to the root directory
+// that it should be loaded from.
+type testingLoader struct {
+        cacheDir string
+
+        virtualEnvZIP string
+
+        pantsWheelPath string
+        shirtWheelPath string
+}
+
+// loadTestEnvironment sets up the test environment for the VirtualEnv tests.
+//
+// This environment includes the acquisition and construction of binary data
+// that will be used to perform the VirtualEnv test suite, namely:
+//
+//      - Building test wheel files from source.
+//      - Downloading the testing VirtualEnv package.
+//
+// This online setup is preferred to actually checking these binary files into
+// Git, as it offers more versatility and doesn't clutter Git with binary junk.
+//
+// To optimize repeated test re-executions, withTestEnvironment will also cache
+// the downloaded artifacts in a cache directory. All artifacts will be verified
+// by their SHA256 hashes, which will be baked into the source here.
+func loadTestEnvironment(ctx context.Context, t *testing.T) (*testingLoader, error) {
+        wd, err := os.Getwd()
+        if err != nil {
+                return nil, errors.Annotate(err).Reason("failed to get working directory").Err()
+        }
+
+        cacheDir := filepath.Join(wd, ".venv_test_cache")
+        if err := filesystem.MakeDirs(cacheDir); err != nil {
+                return nil, errors.Annotate(err).Reason("failed to create cache dir").Err()
+        }
+
+        tl := testingLoader{
+                cacheDir: cacheDir,
+        }
+        return &tl, tl.withCacheLock(t, func() error {
+                return tl.ensureRemoteFilesLocked(ctx, t)
+        })
+}
+
+func (tl *testingLoader) withCacheLock(t *testing.T, fn func() error) error {
+        lockPath := filepath.Join(tl.cacheDir, ".lock")
+        blocker := func() error {
+                t.Logf("Cache [%s] is currently locked; sleeping...", lockPath)
+                time.Sleep(1 * time.Second)
+                return nil
+        }
+        return fslock.WithBlocking(lockPath, blocker, func() error {
+                return fn()
+        })
+}
+
+func (tl *testingLoader) ensureWheels(ctx context.Context, t *testing.T, py *python.Interpreter, tdir string) error {
+        var err error
+        if tl.pantsWheelPath, err = tl.buildWheelLocked(t, py, "pants-1.2-py2.py3-none-any.whl", tdir); err != nil {
+                return err
+        }
+        if tl.shirtWheelPath, err = tl.buildWheelLocked(t, py, "shirt-3.14-py2.py3-none-any.whl", tdir); err != nil {
+                return err
+        }
+        return nil
+}
+
+func (tl *testingLoader) Resolve(c context.Context, root string, packages []*vpython.Spec_Package) error {
+        for _, pkg := range packages {
+                pkg.Version = "resolved"
+        }
+        return nil
+}
+
+func (tl *testingLoader) Ensure(c context.Context, root string, packages []*vpython.Spec_Package) error {
+        for _, pkg := range packages {
+                if err := tl.installPackage(pkg.Path, root); err != nil {
+                        return err
+                }
+        }
+        return nil
+}
+
+func (tl *testingLoader) installPackage(name, root string) error {
+        switch name {
+        case "foo/bar/virtualenv":
+                return unzip(tl.virtualEnvZIP, root)
+        case "foo/bar/shirt":
+                return copyFileIntoDir(tl.shirtWheelPath, root)
+        case "foo/bar/pants":
+                return copyFileIntoDir(tl.pantsWheelPath, root)
+
+        default:
+                return errors.Reason("don't know how to install %(package)q").
+                        D("package", name).
+                        Err()
+        }
+}
+
+func (tl *testingLoader) buildWheelLocked(t *testing.T, py *python.Interpreter, name, outDir string) (string, error) {
+        w, err := wheel.ParseName(name)
+        if err != nil {
+                return "", errors.Annotate(err).Reason("failed to parse wheel name %(name)q").
+                        D("name", name).
+                        Err()
+        }
+
+        outWheelPath := filepath.Join(outDir, w.String())
+        switch _, err := os.Stat(outWheelPath); {
+        case err == nil:
+                t.Logf("Using cached wheel for %q: %s", name, outWheelPath)
+                return outWheelPath, nil
+
+        case os.IsNotExist(err):
+                // Will build a new wheel.
+                break
+
+        default:
+                return "", errors.Annotate(err).Reason("failed to stat wheel path [%(path)s]").
+                        D("path", outWheelPath).
+                        Err()
+        }
+
+        srcDir := filepath.Join(testDataDir, w.Distribution+".src")
+
+        // Create a bootstrap wheel-generating VirtualEnv!
+        cfg := Config{
+                MaxHashLen: 1, // Only going to be 1 enviroment.
+                BaseDir:    filepath.Join(outDir, "wheel-vpython"),
+                Python:     py.Python,
+                Package: vpython.Spec_Package{
+                        Path:    "foo/bar/virtualenv",
+                        Version: "whatever",
+                },
+                Loader: tl,
+                Spec:   &vpython.Spec{},
+
+                // Testing parameters for this bootstrap wheel-building environment.
+                testPreserveInstallationCapability: true,
+                testLeaveReadWrite:                 true,
+        }
+
+        return outWheelPath, With(context.Background(), cfg, true, func(ctx context.Context, env *Env) error {
+                cmd := env.InterpreterCommand()
+                cmd.WorkDir = srcDir
+                if err := cmd.Run(context.Background(), "setup.py", "bdist_wheel", "--dist-dir", outDir); err != nil {
+                        return errors.Annotate(err).Reason("failed to build wheel").Err()
+                }
+
+                if _, err := os.Stat(outWheelPath); err != nil {
+                        return errors.Annotate(err).Reason("failed to generate wheel").Err()
+                }
+
+                t.Logf("Generated wheel file %q: %s", name, outWheelPath)
+                return nil
+        })
+}
+
+func (tl *testingLoader) ensureRemoteFilesLocked(ctx context.Context, t *testing.T) error {
+MainLoop:
+        for _, rf := range remoteFiles {
+                cachePath := filepath.Join(tl.cacheDir, rf.name)
+
+                // Check if the remote file is already cached.
+                err := getCachedFileLocked(t, cachePath, rf.contentHash)
+                if err == nil {
+                        t.Logf("Remote file [%s] is already cached: [%s]", rf.name, cachePath)
+                        rf.install(tl, cachePath)
+                        continue MainLoop
+                }
+                t.Logf("Remote file [%s] is not cached: %s", rf.name, err)
+
+                // Download from CIPD.
+                if rf.cipdPackage != "" {
+                        err := cacheFromCIPDLocked(ctx, t, cachePath, rf.name, rf.contentHash, rf.cipdPackage, rf.cipdVersion)
+                        if err == nil {
+                                t.Logf("Cached remote file [%s] from CIPD source: [%s]", rf.name, cachePath)
+                                rf.install(tl, cachePath)
+                                continue MainLoop
+                        }
+                        t.Logf("Failed to load from CIPD package %q @%q: %s", rf.cipdPackage, rf.cipdVersion, err)
+                }
+
+                // Download from URL.
+                for _, url := range rf.urls {
+                        err := cacheFromURLLocked(t, cachePath, rf.contentHash, url)
+                        if err == nil {
+                                t.Logf("Cached remote file [%s] from URL [%s]: [%s]", rf.name, url, cachePath)
+                                rf.install(tl, cachePath)
+                                continue MainLoop
+                        }
+                        t.Logf("Failed to load from URL %q: %s", url, err)
+                }
+
+                return errors.Reason("failed to acquire remote file %(name)q").
+                        D("name", rf.name).
+                        Err()
+        }
+
+        return nil
+}
+
+func getCachedFileLocked(t *testing.T, cachePath, hash string) error {
+        return validateHash(t, cachePath, hash, true)
+}
+
+func validateHash(t *testing.T, path, hash string, deleteIfInvalid bool) error {
+        fd, err := os.Open(path)
+        if err != nil {
+                return errors.Annotate(err).Reason("failed to open file").Err()
+        }
+        defer fd.Close()
+
+        h := sha256.New()
+        if _, err := io.Copy(h, fd); err != nil {
+                return errors.Annotate(err).Reason("failed to hash file").Err()
+        }
+
+        if err := hashesEqual(h, hash); err != nil {
+                t.Logf("File [%s] has invalid hash: %s", path, err)
+
+                if deleteIfInvalid {
+                        if err := os.Remove(path); err != nil {
+                                t.Logf("Failed to delete invalid hash file [%s]: %s", path, err)
+                        }
+                }
+                return err
+        }
+
+        return nil
+}
+
+func hashesEqual(h hash.Hash, expected string) error {
+        if v := hex.EncodeToString(h.Sum(nil)); v != expected {
+                return errors.Reason("hash %(actual)q doesn't match expected %(expected)q").
+                        D("actual", v).
+                        D("expected", expected).
+                        Err()
+        }
+        return nil
+}
+
+var testCIPDClientOptions = cipd.ClientOptions{
+        ServiceURL: chromeinfra.CIPDServiceURL,
+        UserAgent:  "vpython venv tests",
+}
+
+func cacheFromCIPDLocked(ctx context.Context, t *testing.T, cachePath, name, hash, pkg, version string) error {
+        return testfs.WithTempDir(t, "vpython_venv_cipd", func(tdir string) error {
+                opts := testCIPDClientOptions
+                opts.Root = tdir
+
+                client, err := cipd.NewClient(opts)
+                if err != nil {
+                        return errors.Annotate(err).Reason("failed to create CIPD client").Err()
+                }
+
+                pin, err := client.ResolveVersion(ctx, pkg, version)
+                if err != nil {
+                        return errors.Annotate(err).Reason("failed to resolve CIPD version for %(pkg)s @%(version)s").
+                                D("pkg", pkg).
+                                D("version", version).
+                                Err()
+                }
+
+                if err := client.FetchAndDeployInstance(ctx, "", pin); err != nil {
+                        return errors.Annotate(err).Reason("failed to fetch/deploy CIPD package").Err()
+                }
+
+                path := filepath.Join(opts.Root, name)
+                if err := validateHash(t, path, hash, false); err != nil {
+                        // Do not export the invalid path.
+                        return err
+                }
+
+                if err := copyFile(path, cachePath, nil); err != nil {
+                        return errors.Annotate(err).Reason("failed to install CIPD package file").Err()
+                }
+
+                return nil
+        })
+}
+
+func cacheFromURLLocked(t *testing.T, cachePath, hash, url string) (err error) {
+        resp, err := http.Get(url)
+        if err != nil {
+                t.Logf("Failed to GET file from URL [%s]: %s", url, err)
+        }
+        defer resp.Body.Close()
+
+        fd, err := os.Create(cachePath)
+        if err != nil {
+                t.Logf("Failed to create output file [%s]: %s", cachePath, err)
+        }
+        defer func() {
+                if closeErr := fd.Close(); closeErr != nil && err == nil {
+                        err = errors.Annotate(closeErr).Reason("failed to close file").Err()
+                }
+        }()
+
+        h := sha256.New()
+        tr := io.TeeReader(resp.Body, h)
+        if _, err := io.Copy(fd, tr); err != nil {
+                return errors.Annotate(err).Reason("failed to download").Err()
+        }
+
+        if err = hashesEqual(h, hash); err != nil {
+                return
+        }
+        return nil
+}
+
+func unzip(src, dst string) error {
+        fd, err := zip.OpenReader(src)
+        if err != nil {
+                return errors.Annotate(err).Reason("failed to open ZIP reader").Err()
+        }
+        defer fd.Close()
+
+        for _, f := range fd.File {
+                path := filepath.Join(dst, filepath.FromSlash(f.Name))
+                fi := f.FileInfo()
+
+                // Unzip this entry.
+                if fi.IsDir() {
+                        if err := os.MkdirAll(path, 0755); err != nil {
+                                return errors.Annotate(err).Reason("failed to mkdir").Err()
+                        }
+                } else {
+                        if err := copyFileOpener(f.Open, path, fi); err != nil {
+                                return err
+                        }
+                }
+        }
+        return nil
+}
+
+func copyFileIntoDir(src, dstDir string) error {
+        return copyFile(src, filepath.Join(dstDir, filepath.Base(src)), nil)
+}
+
+func copyFile(src, dst string, fi os.FileInfo) error {
+        opener := func() (io.ReadCloser, error) { return os.Open(src) }
+        return copyFileOpener(opener, dst, fi)
+}
+
+func copyFileOpener(opener func() (io.ReadCloser, error), dst string, fi os.FileInfo) (err error) {
+        sfd, err := opener()
+        if err != nil {
+                return errors.Annotate(err).Reason("failed to open source").Err()
+        }
+        defer sfd.Close()
+
+        dfd, err := os.Create(dst)
+        if err != nil {
+                return errors.Annotate(err).Reason("failed to create destination").Err()
+        }
+        defer func() {
+                if closeErr := dfd.Close(); closeErr != nil && err == nil {
+                        err = errors.Annotate(closeErr).Reason("failed to close destination").Err()
+                }
+        }()
+
+        if _, err := io.Copy(dfd, sfd); err != nil {
+                return errors.Annotate(err).Reason("failed to copy file").Err()
+        }
+        if fi != nil {
+                if err := os.Chmod(dst, fi.Mode()); err != nil {
+                        return errors.Annotate(err).Reason("failed to chmod").Err()
+                }
+        }
+        return nil
+}