1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
|
# Copyright 2017 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import contextlib
import logging
import os
import tempfile
from devil.utils import reraiser_thread
class Datatype:
HTML = 'text/html'
JSON = 'application/json'
PNG = 'image/png'
TEXT = 'text/plain'
class OutputManager:
def __init__(self):
"""OutputManager Constructor.
This class provides a simple interface to save test output. Subclasses
of this will allow users to save test results in the cloud or locally.
"""
self._allow_upload = False
self._thread_group = None
@contextlib.contextmanager
def ArchivedTempfile(
self, out_filename, out_subdir, datatype=Datatype.TEXT):
"""Archive file contents asynchronously and then deletes file.
Args:
out_filename: Name for saved file.
out_subdir: Directory to save |out_filename| to.
datatype: Datatype of file.
Returns:
An ArchivedFile file. This file will be uploaded async when the context
manager exits. AFTER the context manager exits, you can get the link to
where the file will be stored using the Link() API. You can use typical
file APIs to write and flish the ArchivedFile. You can also use file.name
to get the local filepath to where the underlying file exists. If you do
this, you are responsible of flushing the file before exiting the context
manager.
"""
if not self._allow_upload:
raise Exception('Must run |SetUp| before attempting to upload!')
f = self.CreateArchivedFile(out_filename, out_subdir, datatype)
try:
yield f
finally:
self.ArchiveArchivedFile(f, delete=True)
def CreateArchivedFile(self, out_filename, out_subdir,
datatype=Datatype.TEXT):
"""Returns an instance of ArchivedFile."""
return self._CreateArchivedFile(out_filename, out_subdir, datatype)
def _CreateArchivedFile(self, out_filename, out_subdir, datatype):
raise NotImplementedError
def ArchiveArchivedFile(self, archived_file, delete=False):
"""Archive an ArchivedFile instance and optionally delete it."""
if not isinstance(archived_file, ArchivedFile):
raise Exception('Excepting an instance of ArchivedFile, got %s.' %
type(archived_file))
archived_file.PrepareArchive()
def archive():
try:
archived_file.Archive()
finally:
if delete:
archived_file.Delete()
thread = reraiser_thread.ReraiserThread(func=archive)
thread.start()
self._thread_group.Add(thread)
def SetUp(self):
self._allow_upload = True
self._thread_group = reraiser_thread.ReraiserThreadGroup()
def TearDown(self):
self._allow_upload = False
logging.info('Finishing archiving output.')
self._thread_group.JoinAll()
def __enter__(self):
self.SetUp()
return self
def __exit__(self, _exc_type, _exc_val, _exc_tb):
self.TearDown()
class ArchivedFile:
def __init__(self, out_filename, out_subdir, datatype):
self._out_filename = out_filename
self._out_subdir = out_subdir
self._datatype = datatype
mode = 'w+'
if datatype == Datatype.PNG:
mode = 'w+b'
self._f = tempfile.NamedTemporaryFile(mode=mode, delete=False)
self._ready_to_archive = False
@property
def name(self):
return self._f.name
def fileno(self, *args, **kwargs):
if self._ready_to_archive:
raise Exception('Cannot retrieve the integer file descriptor '
'after archiving has begun!')
return self._f.fileno(*args, **kwargs)
def write(self, *args, **kwargs):
if self._ready_to_archive:
raise Exception('Cannot write to file after archiving has begun!')
self._f.write(*args, **kwargs)
def flush(self, *args, **kwargs):
if self._ready_to_archive:
raise Exception('Cannot flush file after archiving has begun!')
self._f.flush(*args, **kwargs)
def Link(self):
"""Returns location of archived file."""
if not self._ready_to_archive:
raise Exception('Cannot get link to archived file before archiving '
'has begun')
return self._Link()
def _Link(self):
"""Note for when overriding this function.
This function will certainly be called before the file
has finished being archived. Therefore, this needs to be able to know the
exact location of the archived file before it is finished being archived.
"""
raise NotImplementedError
def PrepareArchive(self):
"""Meant to be called synchronously to prepare file for async archiving."""
self.flush()
self._ready_to_archive = True
self._PrepareArchive()
def _PrepareArchive(self):
"""Note for when overriding this function.
This function is needed for things such as computing the location of
content addressed files. This is called after the file is written but
before archiving has begun.
"""
def Archive(self):
"""Archives file."""
if not self._ready_to_archive:
raise Exception('File is not ready to archive. Be sure you are not '
'writing to the file and PrepareArchive has been called')
self._Archive()
def _Archive(self):
raise NotImplementedError
def Delete(self):
"""Deletes the backing file."""
self._f.close()
os.remove(self.name)
|
