Fixed persistance for BayesMemoryBackend (#147)

lib/classifier-reborn/backends/bayes_memory_backend.rb

@@ -6,7 +6,7 @@ class BayesMemoryBackend
     def initialize
       @total_words     = 0
       @total_trainings = 0
-      @category_counts = Hash.new { |h, k| h[k] = {training: 0, word: 0} }
+      @category_counts = {}
       @categories      = {}
     end
 
@@ -19,23 +19,23 @@ def update_total_trainings(diff)
     end
 
     def category_training_count(category)
-      @category_counts[category][:training]
+      category_counts(category)[:training]
     end
 
     def update_category_training_count(category, diff)
-      @category_counts[category][:training] += diff
+      category_counts(category)[:training] += diff
     end
 
     def category_has_trainings?(category)
       @category_counts.key?(category) && category_training_count(category) > 0
     end
 
     def category_word_count(category)
-      @category_counts[category][:word]
+      category_counts(category)[:word]
     end
 
     def update_category_word_count(category, diff)
-      @category_counts[category][:word] += diff
+      category_counts(category)[:word] += diff
     end
 
     def add_category(category)
@@ -65,5 +65,11 @@ def word_in_category?(category, word)
     def reset
       initialize
     end
+
+    private
+
+    def category_counts(category)
+      @category_counts[category] ||= {training: 0, word: 0}
+    end
   end
 end

test/backends/backend_memory_test.rb

@@ -9,4 +9,23 @@ class BackendMemoryTest < Minitest::Test
   def setup
     @backend = ClassifierReborn::BayesMemoryBackend.new
   end
+
+  def test_persistence
+    @backend.update_total_words(10)
+    @backend.update_total_trainings(10)
+    @backend.add_category(:Interesting)
+    @backend.update_category_training_count(:Interesting, 10)
+    @backend.update_category_word_count(:Interesting, 10)
+    @backend.update_category_word_frequency(:Interesting, 'foo', 10)
+
+    binary = Marshal.dump(@backend)
+    loaded_backend = Marshal.load(binary)
+
+    assert_equal @backend.total_words, loaded_backend.total_words
+    assert_equal @backend.total_trainings, loaded_backend.total_trainings
+    assert_equal @backend.category_training_count(:Interesting), loaded_backend.category_training_count(:Interesting)
+    assert_equal @backend.category_word_count(:Interesting), loaded_backend.category_word_count(:Interesting)
+    assert_equal @backend.category_word_frequency(:Interesting, 'foo'),
+                 loaded_backend.category_word_frequency(:Interesting, 'foo')
+  end
 end