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()