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_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_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_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_count_token_occurrences(self): sb = SimpleBayes() result = sb.count_token_occurrences(['hello', 'world', 'hello']) self.assertEqual(result, {'hello': 2, 'world': 1})