import io
import unittest

from wordfreq import WordCount, count_words, run, top_n


class TestWordFrequency(unittest.TestCase):
    def test_count_words(self):
        source = io.StringIO("Hello, world! hello gpu-0 café")
        self.assertEqual(
            count_words(source, True),
            {"hello": 2, "world": 1, "gpu": 1, "0": 1, "café": 1},
        )

    def test_top_n(self):
        counts = {"b": 2, "a": 2, "c": 5, "d": 1}
        self.assertEqual(
            top_n(counts, 3),
            [WordCount("c", 5), WordCount("a", 2), WordCount("b", 2)],
        )

    def test_cli_stdin(self):
        stdout, stderr = io.StringIO(), io.StringIO()
        code = run(["-n", "2"], io.StringIO("x y x z z z"), stdout, stderr)
        self.assertEqual(code, 0, stderr.getvalue())
        self.assertEqual(stdout.getvalue(), "3\tz\n2\tx\n")


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