| OLD | NEW |
| (Empty) |
| 1 # Copyright (c) 2009, Google Inc. All rights reserved. | |
| 2 # | |
| 3 # Redistribution and use in source and binary forms, with or without | |
| 4 # modification, are permitted provided that the following conditions are | |
| 5 # met: | |
| 6 # | |
| 7 # * Redistributions of source code must retain the above copyright | |
| 8 # notice, this list of conditions and the following disclaimer. | |
| 9 # * Redistributions in binary form must reproduce the above | |
| 10 # copyright notice, this list of conditions and the following disclaimer | |
| 11 # in the documentation and/or other materials provided with the | |
| 12 # distribution. | |
| 13 # * Neither the name of Google Inc. nor the names of its | |
| 14 # contributors may be used to endorse or promote products derived from | |
| 15 # this software without specific prior written permission. | |
| 16 # | |
| 17 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS | |
| 18 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT | |
| 19 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR | |
| 20 # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT | |
| 21 # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, | |
| 22 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT | |
| 23 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, | |
| 24 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY | |
| 25 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT | |
| 26 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE | |
| 27 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. | |
| 28 # | |
| 29 # Class for unittest support. Used for capturing stderr/stdout. | |
| 30 | |
| 31 import logging | |
| 32 import sys | |
| 33 | |
| 34 from StringIO import StringIO | |
| 35 | |
| 36 | |
| 37 class OutputCapture(object): | |
| 38 | |
| 39 def __init__(self): | |
| 40 self.saved_outputs = dict() | |
| 41 self._log_level = logging.INFO | |
| 42 | |
| 43 def set_log_level(self, log_level): | |
| 44 self._log_level = log_level | |
| 45 if hasattr(self, '_logs_handler'): | |
| 46 self._logs_handler.setLevel(self._log_level) | |
| 47 | |
| 48 def _capture_output_with_name(self, output_name): | |
| 49 stream = getattr(sys, output_name) | |
| 50 captured_output = StringIO() | |
| 51 self.saved_outputs[output_name] = stream | |
| 52 setattr(sys, output_name, captured_output) | |
| 53 return captured_output | |
| 54 | |
| 55 def _restore_output_with_name(self, output_name): | |
| 56 captured_output = getattr(sys, output_name).getvalue() | |
| 57 setattr(sys, output_name, self.saved_outputs[output_name]) | |
| 58 del self.saved_outputs[output_name] | |
| 59 return captured_output | |
| 60 | |
| 61 def capture_output(self): | |
| 62 self._logs = StringIO() | |
| 63 self._logs_handler = logging.StreamHandler(self._logs) | |
| 64 self._logs_handler.setLevel(self._log_level) | |
| 65 self._logger = logging.getLogger() | |
| 66 self._orig_log_level = self._logger.level | |
| 67 self._logger.addHandler(self._logs_handler) | |
| 68 self._logger.setLevel(min(self._log_level, self._orig_log_level)) | |
| 69 return (self._capture_output_with_name("stdout"), self._capture_output_w
ith_name("stderr")) | |
| 70 | |
| 71 def restore_output(self): | |
| 72 self._logger.removeHandler(self._logs_handler) | |
| 73 self._logger.setLevel(self._orig_log_level) | |
| 74 self._logs_handler.flush() | |
| 75 self._logs.flush() | |
| 76 logs_string = self._logs.getvalue() | |
| 77 delattr(self, '_logs_handler') | |
| 78 delattr(self, '_logs') | |
| 79 return (self._restore_output_with_name("stdout"), self._restore_output_w
ith_name("stderr"), logs_string) | |
| 80 | |
| 81 def assert_outputs(self, testcase, function, args=[], kwargs={}, expected_st
dout="", | |
| 82 expected_stderr="", expected_exception=None, expected_log
s=None): | |
| 83 self.capture_output() | |
| 84 try: | |
| 85 if expected_exception: | |
| 86 return_value = testcase.assertRaises(expected_exception, functio
n, *args, **kwargs) | |
| 87 else: | |
| 88 return_value = function(*args, **kwargs) | |
| 89 finally: | |
| 90 (stdout_string, stderr_string, logs_string) = self.restore_output() | |
| 91 | |
| 92 if hasattr(testcase, 'assertMultiLineEqual'): | |
| 93 testassert = testcase.assertMultiLineEqual | |
| 94 else: | |
| 95 testassert = testcase.assertEqual | |
| 96 | |
| 97 testassert(stdout_string, expected_stdout) | |
| 98 testassert(stderr_string, expected_stderr) | |
| 99 if expected_logs is not None: | |
| 100 testassert(logs_string, expected_logs) | |
| 101 # This is a little strange, but I don't know where else to return this i
nformation. | |
| 102 return return_value | |
| OLD | NEW |