def test_untrain(self):
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
self.assertEqual(sb.tally('foo'), 3)
self.assertEqual(sb.tally('bar'), 0)
sb.untrain('bar', 'for bar baz')
self.assertEqual(sb.tally('foo'), 3)
self.assertEqual(sb.tally('bar'), 0)
sb.untrain('foo', 'hello world')
self.assertEqual(sb.tally('foo'), 1)
def test_untrain(self):
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
self.assertEqual(sb.tally('foo'), 3)
self.assertEqual(sb.tally('bar'), 0)
sb.untrain('bar', 'for bar baz')
self.assertEqual(sb.tally('foo'), 3)
self.assertEqual(sb.tally('bar'), 0)
sb.untrain('foo', 'hello world')
self.assertEqual(sb.tally('foo'), 1)
def test_train_with_existing_category(self, get_category_mock):
cat_mock = mock.MagicMock()
cat_mock.train_token.return_value = None
get_category_mock.return_value = cat_mock
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
get_category_mock.assert_called_once_with('foo')
cat_mock.train_token.assert_any_call('hello', 2)
cat_mock.train_token.assert_any_call('world', 1)
def test_train_with_existing_category(self, get_category_mock):
cat_mock = mock.MagicMock()
cat_mock.train_token.return_value = None
get_category_mock.return_value = cat_mock
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
get_category_mock.assert_called_once_with('foo')
cat_mock.train_token.assert_any_call('hello', 2)
cat_mock.train_token.assert_any_call('world', 1)
def test_train_with_new_category(self, add_category_mock,
get_category_mock):
cat_mock = mock.MagicMock()
cat_mock.train_token.return_value = None
get_category_mock.side_effect = KeyError()
add_category_mock.return_value = cat_mock
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
add_category_mock.assert_called_with('foo')
cat_mock.train_token.assert_any_call('hello', 2)
cat_mock.train_token.assert_any_call('world', 1)
def test_train_with_new_category(
self,
add_category_mock,
get_category_mock
):
cat_mock = mock.MagicMock()
cat_mock.train_token.return_value = None
get_category_mock.side_effect = KeyError()
add_category_mock.return_value = cat_mock
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
add_category_mock.assert_called_with('foo')
cat_mock.train_token.assert_any_call('hello', 2)
cat_mock.train_token.assert_any_call('world', 1)
def test_flush_and_tally(self):
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
self.assertEqual(sb.tally('foo'), 3)
sb.flush()
self.assertEqual(sb.tally('foo'), 0)
def test_flush_and_tally(self):
sb = SimpleBayes()
sb.train('foo', 'hello world hello')
self.assertEqual(sb.tally('foo'), 3)
sb.flush()
self.assertEqual(sb.tally('foo'), 0)
