unittestで標準出力をテスト

標準出力の内容をテストしたいときはdoctestを使うほうが簡単だけど
unittestを使いたかったのでメモ。

with構文で、標準出力をフックしたあと必ず元に戻すクラスを書いてみた。

import sys
from cStringIO import StringIO

class HookStdOut(object):
    def __init__(self):
        self.fp = StringIO()
        self.orig_stdout = sys.stdout
        sys.stdout = self.fp
    def dump(self):
        self.fp.seek(0)
        return self.fp.read()
    def __enter__(self):
        return self
    def __exit__(self, *args):
        sys.stdout = self.orig_stdout

if __name__ == '__main__':
    import unittest

    def saySomething():
        print "blah blah blah"
    
    class TestSaySomething(unittest.TestCase):
        def setUp(self):
            pass
        def testSaySomething(self):
            with HookStdOut() as hook:
                saySomething()
            self.assertEquals(hook.dump(), "blah blah blah\n")    
            print "You hooked sys.stdout, and now sys.stdout is what it was."

    unittest.main()

__exit__はエラー処理もちゃんと書いたほうがいいかもしれない。