dataset/problem_7/test.py

import unittest
from dataset.problem_7.buggy import top_k_frequent


class TestTopKFrequent(unittest.TestCase):
    def test_simple(self):
        self.assertEqual(top_k_frequent([1, 1, 1, 2, 2, 3], 2), [1, 2])

    def test_single(self):
        self.assertEqual(top_k_frequent([4, 4, 5], 1), [4])


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