dataset/problem_1/test.py

import unittest
from dataset.problem_1.buggy import safe_divide


class TestSafeDivide(unittest.TestCase):
    def test_normal(self):
        self.assertEqual(safe_divide(8, 2), 4)

    def test_zero_division(self):
        self.assertEqual(safe_divide(1, 0), float("inf"))

    def test_type_error_should_raise(self):
        with self.assertRaises(TypeError):
            safe_divide("1", 1)


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