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
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
|
import atexit
import codecs
import os
import sys
import time
import unicodedata
from abc import ABC, abstractmethod
from contextlib import ExitStack, contextmanager
from pathlib import Path
from queue import Empty, Queue
from typing import Any, Dict, Iterator, List
from xml.sax.saxutils import XMLGenerator
from xml.sax.xmlreader import AttributesImpl
from colorama import Fore, Style
from junit_xml import TestCase, TestSuite
class AbstractLogger(ABC):
@abstractmethod
def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
pass
@abstractmethod
@contextmanager
def subtest(self, name: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
pass
@abstractmethod
@contextmanager
def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
pass
@abstractmethod
def info(self, *args, **kwargs) -> None: # type: ignore
pass
@abstractmethod
def warning(self, *args, **kwargs) -> None: # type: ignore
pass
@abstractmethod
def error(self, *args, **kwargs) -> None: # type: ignore
pass
@abstractmethod
def log_serial(self, message: str, machine: str) -> None:
pass
@abstractmethod
def print_serial_logs(self, enable: bool) -> None:
pass
class JunitXMLLogger(AbstractLogger):
class TestCaseState:
def __init__(self) -> None:
self.stdout = ""
self.stderr = ""
self.failure = False
def __init__(self, outfile: Path) -> None:
self.tests: dict[str, JunitXMLLogger.TestCaseState] = {
"main": self.TestCaseState()
}
self.currentSubtest = "main"
self.outfile: Path = outfile
self._print_serial_logs = True
atexit.register(self.close)
def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
self.tests[self.currentSubtest].stdout += message + os.linesep
@contextmanager
def subtest(self, name: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
old_test = self.currentSubtest
self.tests.setdefault(name, self.TestCaseState())
self.currentSubtest = name
yield
self.currentSubtest = old_test
@contextmanager
def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
self.log(message)
yield
def info(self, *args, **kwargs) -> None: # type: ignore
self.tests[self.currentSubtest].stdout += args[0] + os.linesep
def warning(self, *args, **kwargs) -> None: # type: ignore
self.tests[self.currentSubtest].stdout += args[0] + os.linesep
def error(self, *args, **kwargs) -> None: # type: ignore
self.tests[self.currentSubtest].stderr += args[0] + os.linesep
self.tests[self.currentSubtest].failure = True
def log_serial(self, message: str, machine: str) -> None:
if not self._print_serial_logs:
return
self.log(f"{machine} # {message}")
def print_serial_logs(self, enable: bool) -> None:
self._print_serial_logs = enable
def close(self) -> None:
with open(self.outfile, "w") as f:
test_cases = []
for name, test_case_state in self.tests.items():
tc = TestCase(
name,
stdout=test_case_state.stdout,
stderr=test_case_state.stderr,
)
if test_case_state.failure:
tc.add_failure_info("test case failed")
test_cases.append(tc)
ts = TestSuite("NixOS integration test", test_cases)
f.write(TestSuite.to_xml_string([ts]))
class CompositeLogger(AbstractLogger):
def __init__(self, logger_list: List[AbstractLogger]) -> None:
self.logger_list = logger_list
def add_logger(self, logger: AbstractLogger) -> None:
self.logger_list.append(logger)
def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
for logger in self.logger_list:
logger.log(message, attributes)
@contextmanager
def subtest(self, name: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
with ExitStack() as stack:
for logger in self.logger_list:
stack.enter_context(logger.subtest(name, attributes))
yield
@contextmanager
def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
with ExitStack() as stack:
for logger in self.logger_list:
stack.enter_context(logger.nested(message, attributes))
yield
def info(self, *args, **kwargs) -> None: # type: ignore
for logger in self.logger_list:
logger.info(*args, **kwargs)
def warning(self, *args, **kwargs) -> None: # type: ignore
for logger in self.logger_list:
logger.warning(*args, **kwargs)
def error(self, *args, **kwargs) -> None: # type: ignore
for logger in self.logger_list:
logger.error(*args, **kwargs)
sys.exit(1)
def print_serial_logs(self, enable: bool) -> None:
for logger in self.logger_list:
logger.print_serial_logs(enable)
def log_serial(self, message: str, machine: str) -> None:
for logger in self.logger_list:
logger.log_serial(message, machine)
class TerminalLogger(AbstractLogger):
def __init__(self) -> None:
self._print_serial_logs = True
def maybe_prefix(self, message: str, attributes: Dict[str, str]) -> str:
if "machine" in attributes:
return f"{attributes['machine']}: {message}"
return message
@staticmethod
def _eprint(*args: object, **kwargs: Any) -> None:
print(*args, file=sys.stderr, **kwargs)
def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
self._eprint(self.maybe_prefix(message, attributes))
@contextmanager
def subtest(self, name: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
with self.nested("subtest: " + name, attributes):
yield
@contextmanager
def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
self._eprint(
self.maybe_prefix(
Style.BRIGHT + Fore.GREEN + message + Style.RESET_ALL, attributes
)
)
tic = time.time()
yield
toc = time.time()
self.log(f"(finished: {message}, in {toc - tic:.2f} seconds)")
def info(self, *args, **kwargs) -> None: # type: ignore
self.log(*args, **kwargs)
def warning(self, *args, **kwargs) -> None: # type: ignore
self.log(*args, **kwargs)
def error(self, *args, **kwargs) -> None: # type: ignore
self.log(*args, **kwargs)
def print_serial_logs(self, enable: bool) -> None:
self._print_serial_logs = enable
def log_serial(self, message: str, machine: str) -> None:
if not self._print_serial_logs:
return
self._eprint(Style.DIM + f"{machine} # {message}" + Style.RESET_ALL)
class XMLLogger(AbstractLogger):
def __init__(self) -> None:
self.logfile = os.environ.get("LOGFILE", "/dev/null")
self.logfile_handle = codecs.open(self.logfile, "wb")
self.xml = XMLGenerator(self.logfile_handle, encoding="utf-8")
self.queue: "Queue[Dict[str, str]]" = Queue()
self._print_serial_logs = True
self.xml.startDocument()
self.xml.startElement("logfile", attrs=AttributesImpl({}))
def close(self) -> None:
self.xml.endElement("logfile")
self.xml.endDocument()
self.logfile_handle.close()
def sanitise(self, message: str) -> str:
return "".join(ch for ch in message if unicodedata.category(ch)[0] != "C")
def maybe_prefix(self, message: str, attributes: Dict[str, str]) -> str:
if "machine" in attributes:
return f"{attributes['machine']}: {message}"
return message
def log_line(self, message: str, attributes: Dict[str, str]) -> None:
self.xml.startElement("line", attrs=AttributesImpl(attributes))
self.xml.characters(message)
self.xml.endElement("line")
def info(self, *args, **kwargs) -> None: # type: ignore
self.log(*args, **kwargs)
def warning(self, *args, **kwargs) -> None: # type: ignore
self.log(*args, **kwargs)
def error(self, *args, **kwargs) -> None: # type: ignore
self.log(*args, **kwargs)
def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
self.drain_log_queue()
self.log_line(message, attributes)
def print_serial_logs(self, enable: bool) -> None:
self._print_serial_logs = enable
def log_serial(self, message: str, machine: str) -> None:
if not self._print_serial_logs:
return
self.enqueue({"msg": message, "machine": machine, "type": "serial"})
def enqueue(self, item: Dict[str, str]) -> None:
self.queue.put(item)
def drain_log_queue(self) -> None:
try:
while True:
item = self.queue.get_nowait()
msg = self.sanitise(item["msg"])
del item["msg"]
self.log_line(msg, item)
except Empty:
pass
@contextmanager
def subtest(self, name: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
with self.nested("subtest: " + name, attributes):
yield
@contextmanager
def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
self.xml.startElement("nest", attrs=AttributesImpl({}))
self.xml.startElement("head", attrs=AttributesImpl(attributes))
self.xml.characters(message)
self.xml.endElement("head")
tic = time.time()
self.drain_log_queue()
yield
self.drain_log_queue()
toc = time.time()
self.log(f"(finished: {message}, in {toc - tic:.2f} seconds)")
self.xml.endElement("nest")
terminal_logger = TerminalLogger()
xml_logger = XMLLogger()
rootlog: CompositeLogger = CompositeLogger([terminal_logger, xml_logger])
|
