depot_tools/tests/gsutil_test.py

#!/usr/bin/env vpython3
# Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Test gsutil.py."""

from __future__ import unicode_literals

import base64
import hashlib
import io
import json
import os
import shutil
import subprocess
import sys
import tempfile
import unittest
import zipfile
import urllib.request

# Add depot_tools to path
THIS_DIR = os.path.dirname(os.path.abspath(__file__))
DEPOT_TOOLS_DIR = os.path.dirname(THIS_DIR)
sys.path.append(DEPOT_TOOLS_DIR)

import gsutil


class TestError(Exception):
    pass


class FakeCall(object):
    def __init__(self):
        self.expectations = []

    def add_expectation(self, *args, **kwargs):
        returns = kwargs.pop('_returns', None)
        self.expectations.append((args, kwargs, returns))

    def __call__(self, *args, **kwargs):
        if not self.expectations:
            raise TestError('Got unexpected\n%s\n%s' % (args, kwargs))
        exp_args, exp_kwargs, exp_returns = self.expectations.pop(0)
        if args != exp_args or kwargs != exp_kwargs:
            message = 'Expected:\n  args: %s\n  kwargs: %s\n' % (exp_args,
                                                                 exp_kwargs)
            message += 'Got:\n  args: %s\n  kwargs: %s\n' % (args, kwargs)
            raise TestError(message)
        return exp_returns


class GsutilUnitTests(unittest.TestCase):
    def setUp(self):
        self.fake = FakeCall()
        self.tempdir = tempfile.mkdtemp()
        self.old_urlopen = getattr(urllib.request, 'urlopen')
        self.old_call = getattr(subprocess, 'call')
        setattr(urllib.request, 'urlopen', self.fake)
        setattr(subprocess, 'call', self.fake)

    def tearDown(self):
        self.assertEqual(self.fake.expectations, [])
        shutil.rmtree(self.tempdir)
        setattr(urllib.request, 'urlopen', self.old_urlopen)
        setattr(subprocess, 'call', self.old_call)

    def test_download_gsutil(self):
        version = gsutil.VERSION
        filename = 'gsutil_%s.zip' % version
        full_filename = os.path.join(self.tempdir, filename)
        fake_file = b'This is gsutil.zip'
        fake_file2 = b'This is other gsutil.zip'
        url = '%s%s' % (gsutil.GSUTIL_URL, filename)
        self.fake.add_expectation(url, _returns=io.BytesIO(fake_file))

        self.assertEqual(gsutil.download_gsutil(version, self.tempdir),
                         full_filename)
        with open(full_filename, 'rb') as f:
            self.assertEqual(fake_file, f.read())

        metadata_url = gsutil.API_URL + filename
        md5_calc = hashlib.md5()
        md5_calc.update(fake_file)
        b64_md5 = base64.b64encode(md5_calc.hexdigest().encode('utf-8'))
        self.fake.add_expectation(metadata_url,
                                  _returns=io.BytesIO(
                                      json.dumps({
                                          'md5Hash':
                                          b64_md5.decode('utf-8')
                                      }).encode('utf-8')))
        self.assertEqual(gsutil.download_gsutil(version, self.tempdir),
                         full_filename)
        with open(full_filename, 'rb') as f:
            self.assertEqual(fake_file, f.read())
        self.assertEqual(self.fake.expectations, [])

        self.fake.add_expectation(
            metadata_url,
            _returns=io.BytesIO(
                json.dumps({
                    'md5Hash':
                    base64.b64encode(b'aaaaaaa').decode('utf-8')  # Bad MD5
                }).encode('utf-8')))
        self.fake.add_expectation(url, _returns=io.BytesIO(fake_file2))
        self.assertEqual(gsutil.download_gsutil(version, self.tempdir),
                         full_filename)
        with open(full_filename, 'rb') as f:
            self.assertEqual(fake_file2, f.read())
        self.assertEqual(self.fake.expectations, [])

    def test_ensure_gsutil_full(self):
        version = gsutil.VERSION
        gsutil_dir = os.path.join(self.tempdir, 'gsutil_%s' % version, 'gsutil')
        gsutil_bin = os.path.join(gsutil_dir, 'gsutil')
        gsutil_flag = os.path.join(gsutil_dir, 'install.flag')
        os.makedirs(gsutil_dir)

        zip_filename = 'gsutil_%s.zip' % version
        url = '%s%s' % (gsutil.GSUTIL_URL, zip_filename)
        _, tempzip = tempfile.mkstemp()
        fake_gsutil = 'Fake gsutil'
        with zipfile.ZipFile(tempzip, 'w') as zf:
            zf.writestr('gsutil/gsutil', fake_gsutil)
        with open(tempzip, 'rb') as f:
            self.fake.add_expectation(url, _returns=io.BytesIO(f.read()))

        # This should write the gsutil_bin with 'Fake gsutil'
        gsutil.ensure_gsutil(version, self.tempdir, False)
        self.assertTrue(os.path.exists(gsutil_bin))
        with open(gsutil_bin, 'r') as f:
            self.assertEqual(f.read(), fake_gsutil)
        self.assertTrue(os.path.exists(gsutil_flag))
        self.assertEqual(self.fake.expectations, [])

    def test_ensure_gsutil_short(self):
        version = gsutil.VERSION
        gsutil_dir = os.path.join(self.tempdir, 'gsutil_%s' % version, 'gsutil')
        gsutil_bin = os.path.join(gsutil_dir, 'gsutil')
        gsutil_flag = os.path.join(gsutil_dir, 'install.flag')
        os.makedirs(gsutil_dir)

        with open(gsutil_bin, 'w') as f:
            f.write('Foobar')
        with open(gsutil_flag, 'w') as f:
            f.write('Barbaz')
        self.assertEqual(gsutil.ensure_gsutil(version, self.tempdir, False),
                         gsutil_bin)


if __name__ == '__main__':
    unittest.main()