construction-responsibility-assignment / tests /test_classification_utils.py
hbaltuntas's picture
Create tests/test_classification_utils.py
c2f880c verified
Raw
History Blame Contribute Delete
1.56 kB
import unittest
import pandas as pd
from classification_utils import class_position, combined_confidence, map_category_label
class ClassificationUtilsTests(unittest.TestCase):
def test_known_numeric_categories_are_mapped(self):
self.assertEqual(map_category_label(0), "Contract")
self.assertEqual(map_category_label("15"), "Construction Works")
self.assertEqual(map_category_label(6.0), "Occupational Health and Safety")
def test_free_text_and_unknown_categories_are_preserved(self):
self.assertEqual(map_category_label("risk"), "risk")
self.assertEqual(map_category_label(99), "99")
self.assertEqual(map_category_label(None), "")
def test_combined_confidence_is_joint_probability(self):
party = pd.Series([0.4, 0.9, None, 1.2])
stakeholder = pd.Series([0.95, 0.8, 0.7, 1.0])
result = combined_confidence(party, stakeholder)
expected = [0.38, 0.72, 0.0, 1.0]
for actual, wanted in zip(result, expected):
self.assertAlmostEqual(actual, wanted)
def test_low_confidence_in_either_stage_keeps_joint_confidence_low(self):
party = pd.Series([0.95])
stakeholder = pd.Series([0.4])
self.assertLess(combined_confidence(party, stakeholder).iloc[0], 0.6)
def test_probability_column_uses_class_order_not_label_value(self):
classes = [10, 20, 30]
self.assertEqual(class_position(classes, 20), 1)
self.assertIsNone(class_position(classes, 99))
if __name__ == "__main__":
unittest.main()