dataset/problem_6/test.py

import unittest
from dataset.problem_6.buggy import count_unique_words


class TestCountUniqueWords(unittest.TestCase):
    def test_duplicates(self):
        self.assertEqual(count_unique_words("Hi hi HI"), 1)

    def test_punctuation(self):
        self.assertEqual(count_unique_words("Hello, world! Hello."), 2)


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