dataset/problem_8/test.py

import unittest
from dataset.problem_8.buggy import flatten_one_level


class TestFlattenOneLevel(unittest.TestCase):
    def test_only_one_level(self):
        data = [1, [2, 3], [4, [5, 6]]]
        self.assertEqual(flatten_one_level(data), [1, 2, 3, 4, [5, 6]])


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