Fixing a critical bug that caused scoring not to work after loading d…

…ata from cache
hickeroar · Apr 22, 2015 · cf1ee8f · cf1ee8f
1 parent 041de91
commit cf1ee8f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/simplebayes/__init__.py b/simplebayes/__init__.py
@@ -302,4 +302,7 @@ def cache_train(self):
 
         self.categories = categories
 
+        # Updating our per-category overall probabilities
+        self.calculate_category_probability()
+
         return True
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -234,10 +234,11 @@ def test_score_with_zero_bayes_denon(self, get_categories_mock):
         cat2_mock.get_token_count.assert_any_call('world')
         cat2_mock.get_tally.assert_called_once_with()
 
+    @mock.patch.object(SimpleBayes, 'calculate_category_probability')
     @mock.patch.object(builtins, 'open')
     @mock.patch.object(pickle, 'load')
     @mock.patch.object(os.path, 'exists')
-    def test_cache_train(self, exists_mock, load_mock, open_mock):
+    def test_cache_train(self, exists_mock, load_mock, open_mock, calc_mock):
         categories = BayesCategories()
         categories.categories = {'foo': 'bar'}
 
@@ -251,6 +252,7 @@ def test_cache_train(self, exists_mock, load_mock, open_mock):
         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)