python-startup/test_python_startup.py

import sys

import unittest

from python_startup import PromptCounter


class RecordStdout(object):
    def __init__(self):
        self.value = None

    def write(self, value):
        self.value = value


class TestPromptCounter(unittest.TestCase):
    def setUp(self):
        sys.stdout = RecordStdout()
        self.prompt = PromptCounter('in {} ', 'out {} ')
        self.assertEqual(1, self.prompt.count)

    def tearDown(self):
        sys.stdout = sys.__stdout__

    def test_call_does_not_increment_count(self):
        for x in range(10):
            self.prompt('test')
            self.assertEqual("out 0 'test'\n", sys.stdout.value)
        self.assertEqual(1, self.prompt.count)

    def test_str_does_increment_count(self):
        for x in range(10):
            out = str(self.prompt)
            self.assertEqual('in {} '.format(x + 1), out)
        self.assertEqual(11, self.prompt.count)

    def test_call_and_str(self):
        for x in range(10):
            out = str(self.prompt)
            self.prompt(x)
            self.assertEqual('in {} '.format(x + 1), out)
            self.assertEqual("out {} {}\n".format(x + 1, x), sys.stdout.value)

    def test_call_with_none(self):
        self.prompt(None)
        self.assertIsNone(sys.stdout.value)


if __name__ == '__main__':
    unittest.main()