refinement and test for batch mode on feature_topics

app/models/category_featured_topic.rb

@@ -3,36 +3,43 @@ class CategoryFeaturedTopic < ActiveRecord::Base
   belongs_to :topic
 
   NEXT_CATEGORY_ID_KEY = 'category-featured-topic:next-category-id'.freeze
-  BATCH_SIZE = 100
+  DEFAULT_BATCH_SIZE = 100
 
   # Populates the category featured topics.
-  def self.feature_topics(batched: false)
+  def self.feature_topics(batched: false, batch_size: nil)
     current = {}
     CategoryFeaturedTopic.select(:topic_id, :category_id).order(:rank).each do |f|
       (current[f.category_id] ||= []) << f.topic_id
     end
 
-    next_category_id = batched ? ($redis.get(NEXT_CATEGORY_ID_KEY) || 0) : 0
+    batch_size ||= DEFAULT_BATCH_SIZE
+
+    next_category_id = batched ? $redis.get(NEXT_CATEGORY_ID_KEY).to_i : 0
 
     categories = Category.select(:id, :topic_id, :num_featured_topics)
       .where('id >= ?', next_category_id)
       .order('id ASC')
-      .limit(BATCH_SIZE)
+      .limit(batch_size)
+      .to_a
 
     if batched
-      if categories.count == BATCH_SIZE
+      if categories.length == batch_size
         next_id = Category.where('id > ?', categories.last.id).order('id asc').limit(1).pluck(:id)[0]
-        next_id ? $redis.setex(NEXT_CATEGORY_ID_KEY, 1.day, next_id) : $redis.del(NEXT_CATEGORY_ID_KEY)
+        next_id ? $redis.setex(NEXT_CATEGORY_ID_KEY, 1.day, next_id) : clear_batch!
       else
-        $redis.del(NEXT_CATEGORY_ID_KEY)
+        clear_batch!
       end
     end
 
-    categories.find_each do |c|
+    categories.each do |c|
       CategoryFeaturedTopic.feature_topics_for(c, current[c.id] || [])
     end
   end
 
+  def self.clear_batch!
+    $redis.del(NEXT_CATEGORY_ID_KEY)
+  end
+
   def self.feature_topics_for(c, existing = nil)
     return if c.blank?
 

spec/models/category_featured_topic_spec.rb

@@ -10,6 +10,25 @@ describe CategoryFeaturedTopic do
     let(:category)       { Fabricate(:category) }
     let!(:category_post) { PostCreator.create(user, raw: "I put this post in the category", title: "categorize THIS", category: category.id) }
 
+    it "works in batched mode" do
+      category2 = Fabricate(:category)
+      post2 = create_post(category: category2.id)
+
+      CategoryFeaturedTopic.destroy_all
+      CategoryFeaturedTopic.clear_batch!
+
+      size = Category.order(:id).where('id < ?', category.id).count + 1
+
+      CategoryFeaturedTopic.feature_topics(batched: true, batch_size: size)
+
+      expect(CategoryFeaturedTopic.where(topic_id: category_post.topic_id).count).to eq(1)
+      expect(CategoryFeaturedTopic.where(topic_id: post2.topic_id).count).to eq(0)
+
+      CategoryFeaturedTopic.feature_topics(batched: true, batch_size: size)
+
+      expect(CategoryFeaturedTopic.where(topic_id: post2.topic_id).count).to eq(1)
+    end
+
     it "should feature topics for a secure category" do
 
       # so much dancing, I am thinking fixures make sense here.