def test_score(self, get_categories_mock):
cat1_mock = mock.MagicMock()
cat1_mock.get_token_count.return_value = 2
cat1_mock.get_tally.return_value = 8
cat2_mock = mock.MagicMock()
cat2_mock.get_token_count.return_value = 4
cat2_mock.get_tally.return_value = 32
get_categories_mock.return_value = {'foo': cat1_mock, 'bar': cat2_mock}
sb = SimpleBayes()
sb.calculate_category_probability()
result = sb.score('hello world')
self.assertEqual({
'foo': 0.22222222222222224,
'bar': 1.777777777777778
}, result)
assert 3 == get_categories_mock.call_count, \
get_categories_mock.call_count
cat1_mock.get_token_count.assert_any_call('hello')
cat1_mock.get_token_count.assert_any_call('world')
cat1_mock.get_tally.assert_called_once_with()
cat2_mock.get_token_count.assert_any_call('hello')
cat2_mock.get_token_count.assert_any_call('world')
cat2_mock.get_tally.assert_called_once_with()
def test_score(self, get_categories_mock):
cat1_mock = mock.MagicMock()
cat1_mock.get_token_count.return_value = 2
cat1_mock.get_tally.return_value = 8
cat2_mock = mock.MagicMock()
cat2_mock.get_token_count.return_value = 4
cat2_mock.get_tally.return_value = 32
get_categories_mock.return_value = {
'foo': cat1_mock,
'bar': cat2_mock
}
sb = SimpleBayes()
sb.calculate_category_probability()
result = sb.score('hello world')
self.assertEqual(
{
'foo': 0.22222222222222224,
'bar': 1.777777777777778
},
result
)
assert 3 == get_categories_mock.call_count, \
get_categories_mock.call_count
cat1_mock.get_token_count.assert_any_call('hello')
cat1_mock.get_token_count.assert_any_call('world')
cat1_mock.get_tally.assert_called_once_with()
cat2_mock.get_token_count.assert_any_call('hello')
cat2_mock.get_token_count.assert_any_call('world')
cat2_mock.get_tally.assert_called_once_with()
def test_cache_train_with_no_file(self, exists_mock):
exists_mock.return_value = False
sb = SimpleBayes()
result = sb.cache_train()
exists_mock.assert_called_once_with('/tmp/_simplebayes.pickle')
self.assertFalse(result)
def test_cache_train_with_no_file(self, exists_mock):
exists_mock.return_value = False
sb = SimpleBayes()
result = sb.cache_train()
exists_mock.assert_called_once_with('/tmp/_simplebayes.pickle')
self.assertFalse(result)
def test_classify_without_categories(self, get_categories_mock):
get_categories_mock.return_value = {}
sb = SimpleBayes()
result = sb.classify('hello world')
self.assertIsNone(result)
assert 2 == get_categories_mock.call_count, \
get_categories_mock.call_count
def test_classify_without_categories(self, get_categories_mock):
get_categories_mock.return_value = {}
sb = SimpleBayes()
result = sb.classify('hello world')
self.assertIsNone(result)
assert 2 == get_categories_mock.call_count, \
get_categories_mock.call_count
def test_count_token_occurrences(self):
sb = SimpleBayes()
result = sb.count_token_occurrences(['hello', 'world', 'hello'])
self.assertEqual(
result,
{
'hello': 2,
'world': 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_persist_cache(self, dump_mock, open_mock):
open_mock.return_value = 'opened'
categories = BayesCategories()
categories.categories = {'foo': 'bar'}
sb = SimpleBayes()
sb.cache_path = '/tmp/'
sb.categories = categories
sb.cache_persist()
open_mock.assert_called_once_with('/tmp/_simplebayes.pickle', 'wb')
dump_mock.assert_called_once_with(categories, 'opened')
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_classify_with_empty_category(self, get_categories_mock):
cat_mock = mock.MagicMock()
cat_mock.get_tally.return_value = 0
cat_mock.get_token_count.return_value = 0
get_categories_mock.return_value = {'foo': cat_mock}
sb = SimpleBayes()
sb.calculate_category_probability()
result = sb.classify('hello world')
self.assertIsNone(result)
assert 3 == get_categories_mock.call_count, \
get_categories_mock.call_count
cat_mock.get_tally.assert_called_once_with()
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_classify_with_empty_category(self, get_categories_mock):
cat_mock = mock.MagicMock()
cat_mock.get_tally.return_value = 0
cat_mock.get_token_count.return_value = 0
get_categories_mock.return_value = {
'foo': cat_mock
}
sb = SimpleBayes()
sb.calculate_category_probability()
result = sb.classify('hello world')
self.assertIsNone(result)
assert 3 == get_categories_mock.call_count, \
get_categories_mock.call_count
cat_mock.get_tally.assert_called_once_with()
def test_cache_train(self, exists_mock, load_mock, open_mock, calc_mock):
categories = BayesCategories()
categories.categories = {'foo': 'bar'}
load_mock.return_value = categories
open_mock.return_value = 'opened'
exists_mock.return_value = True
sb = SimpleBayes(cache_path='foo')
sb.cache_train()
exists_mock.assert_called_once_with('foo/_simplebayes.pickle')
open_mock.assert_called_once_with('foo/_simplebayes.pickle', 'rb')
load_mock.assert_called_once_with('opened')
calc_mock.assert_called_once_with()
self.assertEqual(sb.categories, categories)
def test_cache_train(self, exists_mock, load_mock, open_mock, calc_mock):
categories = BayesCategories()
categories.categories = {'foo': 'bar'}
load_mock.return_value = categories
open_mock.return_value = 'opened'
exists_mock.return_value = True
sb = SimpleBayes(cache_path='foo')
sb.cache_train()
exists_mock.assert_called_once_with('foo/_simplebayes.pickle')
open_mock.assert_called_once_with('foo/_simplebayes.pickle', 'rb')
load_mock.assert_called_once_with('opened')
calc_mock.assert_called_once_with()
self.assertEqual(sb.categories, categories)
def test_persist_cache(self, dump_mock, open_mock):
open_mock.return_value = 'opened'
categories = BayesCategories()
categories.categories = {'foo': 'bar'}
sb = SimpleBayes()
sb.cache_path = '/tmp/'
sb.categories = categories
sb.cache_persist()
open_mock.assert_called_once_with('/tmp/_simplebayes.pickle', 'wb')
dump_mock.assert_called_once_with(categories, 'opened')
def test_tokenizer(self):
sb = SimpleBayes()
result = sb.tokenizer('hello world')
self.assertEqual(result, ['hello', 'world'])
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_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_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_count_token_occurrences(self):
sb = SimpleBayes()
result = sb.count_token_occurrences(['hello', 'world', 'hello'])
self.assertEqual(result, {'hello': 2, 'world': 1})
def test_tokenizer(self):
sb = SimpleBayes()
result = sb.tokenizer('hello world')
self.assertEqual(result, ['hello', 'world'])
