import unittest
from nltkma.classify.naivebayes import NaiveBayesClassifier


class NaiveBayesClassifierTest(unittest.TestCase):
    def test_simple(self):
        training_features = [
            ({'nice': True, 'good': True}, 'positive'),
            ({'bad': True, 'mean': True}, 'negative'),
        ]

        classifier = NaiveBayesClassifier.train(training_features)

        result = classifier.prob_classify({'nice': True})
        self.assertTrue(result.prob('positive') > result.prob('negative'))
        self.assertEqual(result.max(), 'positive')

        result = classifier.prob_classify({'bad': True})
        self.assertTrue(result.prob('positive') < result.prob('negative'))
        self.assertEqual(result.max(), 'negative')
