dataset/problem_16/test.py

import unittest
from dataset.problem_16.buggy import top_label


class TestTopLabel(unittest.TestCase):
    def test_select_highest(self):
        scores = {"cat": 0.2, "dog": 0.7, "bird": 0.1}
        self.assertEqual(top_label(scores), "dog")


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