def test_max_label_zero_confidence(self): """ Test that if there *are* classifications but the maximum conf is 0 that at least one label is returned. """ # Mock classification element instance e = mock.MagicMock(spec_set=ClassificationElement) e.get_classification.return_value = {'a': 0.0} ClassificationElement.max_label(e)
def test_max_label(self): """ Test that max_label correctly returns the label key associated with the greatest associated confidence value. """ # Mock classification element instance #: :type: ClassificationElement e = mock.MagicMock(spec_set=ClassificationElement) e.get_classification.return_value = {1: 0, 2: 1, 3: 0.5} expected_label = 2 actual_label = ClassificationElement.max_label(e) self.assertEqual(actual_label, expected_label)