Refactor Inferx::Category#train and Inferx::Category#untrain

lib/inferx/category.rb

@@ -57,18 +57,18 @@ def get(word)
     #
     # @param [Array<String>] words an array of words
     def train(words)
+      return if words.empty?
+
+      increase = words.size
+      words = collect(words)
+
       @redis.pipelined do
-        increase = collect(words).inject(0) do |count, pair|
-          zincrby(pair[1], pair[0])
-          count + pair[1]
-        end
-
-        if increase > 0
-          hincrby(name, increase)
-          @redis.save unless manual?
-          @size += increase
-        end
+        words.each { |word, count| zincrby(count, word) }
+        hincrby(name, increase)
+        @redis.save unless manual?
       end
+
+      @size += increase
     end
 
     # Prepare to enhance the training data. Use for high performance.
@@ -83,27 +83,29 @@ def ready_to_train(&process)
     #
     # @param [Array<String>] words an array of words
     def untrain(words)
-      decrease = 0
+      return if words.empty?
 
-      values = @redis.pipelined do
-        decrease = collect(words).inject(0) do |count, pair|
-          zincrby(-pair[1], pair[0])
-          count + pair[1]
-        end
+      decrease = words.size
+      words = collect(words)
 
-        zremrangebyscore('-inf', 0)
+      scores = @redis.pipelined do
+        words.each { |word, count| zincrby(-count, word) }
       end
 
-      values[0..-2].each do |score|
+      scores.each do |score|
         score = score.to_i
         decrease += score if score < 0
       end
 
-      if decrease > 0
+      return unless decrease > 0
+
+      @redis.pipelined do
+        zremrangebyscore('-inf', 0)
         hincrby(name, -decrease)
         @redis.save unless manual?
-        @size -= decrease
       end
+
+      @size -= decrease
     end
 
     # Prepare to attenuate the training data giving words.

spec/inferx/category_spec.rb

@@ -138,9 +138,14 @@
 describe Inferx::Category, '#untrain' do
   it 'calls Redis#zincrby, Redis#zremrangebyscore and Redis#hincrby' do
     redis = redis_stub do |s|
+      s.stub!(:pipelined).and_return do |&block|
+        block.call
+        %w(3 -2 1)
+      end
+
       s.should_receive(:zincrby).with('inferx:categories:red', -2, 'apple')
       s.should_receive(:zincrby).with('inferx:categories:red', -3, 'strawberry')
-      s.should_receive(:zremrangebyscore).with('inferx:categories:red', '-inf', 0).and_return(%w(3 -2 1))
+      s.should_receive(:zremrangebyscore).with('inferx:categories:red', '-inf', 0)
       s.should_receive(:hincrby).with('inferx:categories', :red, -3)
       s.should_receive(:save)
     end
@@ -151,8 +156,13 @@
 
   it 'decreases the size attribute' do
     redis = redis_stub do |s|
+      s.stub!(:pipelined).and_return do |&block|
+        block.call
+        %w(3 -2 1)
+      end
+
       s.stub!(:zincrby)
-      s.stub!(:zremrangebyscore).and_return(%w(3 -2 1))
+      s.stub!(:zremrangebyscore)
       s.stub!(:hincrby)
     end
 
@@ -162,10 +172,15 @@
   end
 
   context 'with no update' do
-    it 'does not call Redis#hincrby' do
+    it 'does not call Redis#zremrangebyscore and Redis#hincrby' do
       redis = redis_stub do |s|
+        s.stub!(:pipelined).and_return do |&block|
+          block.call
+          %w(-2 -3 2)
+        end
+
         s.stub!(:zincrby)
-        s.stub!(:zremrangebyscore).and_return(%w(-2 -3 2))
+        s.should_not_receive(:zremrangebyscore)
         s.should_not_receive(:hincrby)
         s.should_not_receive(:save)
       end
@@ -178,8 +193,13 @@
   context 'with manual save' do
     it 'does not call Redis#save' do
       redis = redis_stub do |s|
+        s.stub!(:pipelined).and_return do |&block|
+          block.call
+          %w(3 -2 1)
+        end
+
         s.stub!(:zincrby)
-        s.stub!(:zremrangebyscore).and_return(%w(3 -2 1))
+        s.stub!(:zremrangebyscore)
         s.stub!(:hincrby)
         s.should_not_receive(:save)
       end