← Back to branch summary

~le-chi-thu/lava-test/add-verbose-argument

« back to all changes in this revision

Viewing changes to lava_test/utils.py

  • Committer: Paul Larson
  • Date: 2011-09-22 17:58:01 UTC
  • mfrom: (92.1.4 using-lava-tool-2)
  • Revision ID: paul.larson@canonical.com-20110922175801-l99rky1xxsr6k3fn
Big refactoring branch to make lava-test use lava-tool.  Thanks to
Zygmunt and ChiThu!

Show diffs side-by-side

added added

removed removed

Lines of Context:
lava_test/utils.py
1
 
# Copyright (c) 2010 Linaro
 
1
# Copyright (c) 2010, 2011 Linaro
2
2
#
3
3
# This program is free software: you can redistribute it and/or modify
4
4
# it under the terms of the GNU General Public License as published by
13
13
# You should have received a copy of the GNU General Public License
14
14
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
15
15
 
 
16
import contextlib
 
17
import hashlib
 
18
import logging
16
19
import os
17
20
import shutil
18
 
import subprocess
19
 
import sys
20
21
import urllib2
21
22
import urlparse
22
23
 
25
26
_fake_machine = None
26
27
 
27
28
 
28
 
class Tee(file):
29
 
    """ A file-like object that optionally mimics tee functionality.
30
 
 
31
 
    By default, output will go to both stdout and the file specified.
32
 
    Optionally, quiet=True can be used to mute the output to stdout.
33
 
    """
34
 
    def __init__(self, *args, **kwargs):
35
 
        try:
36
 
            self.quiet = kwargs.pop('quiet')
37
 
        except KeyError:
38
 
            self.quiet = False
39
 
        super(Tee, self).__init__(*args, **kwargs)
40
 
 
41
 
    def write(self, data):
42
 
        super(Tee, self).write(data)
43
 
        if self.quiet is False:
44
 
            sys.stdout.write(data)
45
 
 
46
 
 
47
29
def geturl(url, path=""):
48
30
    urlpath = urlparse.urlsplit(url).path
49
31
    filename = os.path.basename(urlpath)
75
57
    if _fake_paths is not None:
76
58
        if path in _fake_paths:
77
59
            path = _fake_paths[path]
78
 
    with open(path) as fd:
79
 
        data = fd.read()
80
 
    return data
 
60
    with open(path, 'rb') as stream:
 
61
        return stream.read()
81
62
 
82
63
 
83
64
def fake_file(path, data=None, newpath=None):
119
100
    _fake_machine = None
120
101
 
121
102
 
122
 
def run_and_log(cmd, fd, quiet=False):
123
 
    """
124
 
    Run a command and log the output to fd
125
 
    """
126
 
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
127
 
        stderr=subprocess.STDOUT, shell=True)
128
 
    while proc.returncode == None:
129
 
        proc.poll()
130
 
        data = proc.stdout.readline()
131
 
        fd.write(data)
132
 
        if quiet is False:
133
 
            sys.stdout.write(data)
134
 
    return proc.returncode
135
 
 
136
 
 
137
103
def get_machine_type():
138
104
    """
139
105
    Return the machine type
142
108
    if _fake_machine is None:
143
109
        return os.uname()[-1]
144
110
    return _fake_machine
 
111
 
 
112
 
 
113
def mkdir_p(dirname):
 
114
    if not os.path.exists(dirname):
 
115
        os.makedirs(dirname)
 
116
 
 
117
 
 
118
@contextlib.contextmanager
 
119
def changed_directory(dirname, make_if_needed=True):
 
120
    """
 
121
    A context manager for running a piece of code in another
 
122
    directory. The directory is created if needed (by default, can
 
123
    be changed with make_if_needed).
 
124
    """
 
125
    orig_dir = os.getcwd()
 
126
    if make_if_needed:
 
127
        mkdir_p(dirname)
 
128
    logging.info("Changing directory to %r", dirname)
 
129
    os.chdir(dirname)
 
130
    try:
 
131
        yield
 
132
    finally:
 
133
        logging.info("Changing directory to %r", orig_dir)
 
134
        os.chdir(orig_dir)
 
135
 
 
136
 
 
137
def merge_dict(merge_into, merge_from):
 
138
    """
 
139
    Merge two dictionaries recursively:
 
140
 
 
141
        1) Simple values are overwritten with a logging.warning() message
 
142
        2) Lists are appended
 
143
        3) Dictionaries are merged recursively
 
144
    """
 
145
    assert isinstance(merge_into, dict)
 
146
    assert isinstance(merge_from, dict)
 
147
    for key in merge_from.iterkeys():
 
148
        if key in merge_into:
 
149
            if (isinstance(merge_from[key], dict)
 
150
                and isinstance(merge_into[key], dict)):
 
151
                merge_dict(merge_into[key], merge_from[key])
 
152
            elif (isinstance(merge_from[key], list)
 
153
                  and isinstance(merge_into[key], list)):
 
154
                merge_into[key].extend(merge_from[key])
 
155
            else:
 
156
                logging.warning(
 
157
                    "Overwriting existing value of %r:"
 
158
                    "%r overwritten with %r",
 
159
                    key, merge_into[key], merge_from[key])
 
160
                merge_into[key] = merge_from[key]
 
161
        else:
 
162
            merge_into[key] = merge_from[key]
 
163
 
 
164
 
 
165
class Cache(object):
 
166
    """
 
167
    Simple open-cached-URL class
 
168
    """
 
169
 
 
170
    _instance = None
 
171
 
 
172
    def __init__(self):
 
173
        home = os.environ.get('HOME', '/')
 
174
        basecache = os.environ.get('XDG_CACHE_HOME',
 
175
                     os.path.join(home, '.cache'))
 
176
        self.cache_dir = os.path.join(basecache, 'lava_test')
 
177
 
 
178
    @classmethod
 
179
    def get_instance(cls):
 
180
        if cls._instance is None:
 
181
            cls._instance = cls()
 
182
        return cls._instance
 
183
 
 
184
    def open_cached(self, key, mode="r"):
 
185
        """
 
186
        Acts like open() but the pathname is relative to the
 
187
        lava_test-specific cache directory.
 
188
        """
 
189
        if "w" in mode and not os.path.exists(self.cache_dir):
 
190
            os.makedirs(self.cache_dir)
 
191
        if os.path.isabs(key):
 
192
            raise ValueError("key cannot be an absolute path")
 
193
        try:
 
194
            stream = open(os.path.join(self.cache_dir, key), mode)
 
195
            yield stream
 
196
        finally:
 
197
            stream.close()
 
198
 
 
199
    def _key_for_url(self, url):
 
200
        return hashlib.sha1(url).hexdigest()
 
201
 
 
202
    def _refresh_url_cache(self, key, url):
 
203
        with contextlib.nested(
 
204
            contextlib.closing(urllib2.urlopen(url)),
 
205
            self.open_cached(key, "wb")) as (in_stream, out_stream):
 
206
            out_stream.write(in_stream.read())
 
207
 
 
208
    @contextlib.contextmanager
 
209
    def open_cached_url(self, url):
 
210
        """
 
211
        Like urlopen.open() but the content may be cached.
 
212
        """
 
213
        # Do not cache local files, this is not what users would expect
 
214
 
 
215
        # workaround - not using cache at all.
 
216
        # TODO: fix this and use the cache
 
217
        # if url.startswith("file://"):
 
218
        if True:
 
219
            stream = urllib2.urlopen(url)
 
220
        else:
 
221
            key = self._key_for_url(url)
 
222
            try:
 
223
                stream = self.open_cached(key, "rb")
 
224
            except IOError:
 
225
                self._refresh_url_cache(key, url)
 
226
                stream = self.open_cached(key, "rb")
 
227
        try:
 
228
            yield stream
 
229
        finally:
 
230
            stream.close()