Replace results count with something that works and added tests

lib/msf/base/simple/job_state_tracker.rb

@@ -5,11 +5,10 @@ class JobStateTracker
   include MonitorMixin
 
   def initialize(result_ttl=nil)
-    self.results_size = 0
     self.ready = Set.new
     self.running = Set.new
     # Can be expanded upon later to allow the option of a MemCacheStore being backed by redis for example
-    self.results = ActiveSupport::Cache::MemoryStore.new(expires_in: result_ttl || 5.minutes)
+    self.results = ResultsMemoryStore.new(expires_in: result_ttl || 5.minutes)
   end
 
   def waiting(id)
@@ -24,7 +23,7 @@ def start(id)
   def completed(id, result, ttl=nil)
     begin
       # ttl of nil means it will take the default expiry time
-      self.results_size += 1 if results.write(id, {result: result}, ttl)
+      results.write(id, {result: result}, ttl)
     ensure
       running.delete(id)
     end
@@ -33,7 +32,7 @@ def completed(id, result, ttl=nil)
   def failed(id, error, ttl=nil)
     begin
       # ttl of nil means it will take the default expiry time
-      self.results_size += 1 if results.write(id, {error: error.to_s}, ttl)
+      results.write(id, {error: error.to_s}, ttl)
     ensure
       running.delete(id)
     end
@@ -56,9 +55,11 @@ def result(id)
   end
 
   def delete(id)
-    result_deleted = results.delete(id)
-    self.results_size -= 1 if result_deleted
-    result_deleted
+    results.delete(id)
+  end
+
+  def results_size
+    results.size
   end
 
   def waiting_size
@@ -70,11 +71,14 @@ def running_size
   end
 
   alias :ack :delete
-
-  attr_accessor :results_size
-
+
   private
 
-  attr_writer :results_size
   attr_accessor :ready, :running, :results
+
+  class ResultsMemoryStore < ActiveSupport::Cache::MemoryStore
+    def size
+      @data.size
+    end
+  end
 end

spec/lib/msf/base/simple/job_state_tracker_spec.rb

@@ -41,14 +41,17 @@
 
       it 'should show as waiting' do
         expect(job_state_tracker).to be_waiting(job_id)
+        expect(job_state_tracker.waiting_size).to be(1)
       end
 
       it 'should not show as running' do
         expect(job_state_tracker).not_to be_running(job_id)
+        expect(job_state_tracker.running_size).to be(0)
       end
 
       it 'should not show as finished' do
         expect(job_state_tracker).not_to be_finished(job_id)
+        expect(job_state_tracker.results_size).to be(0)
       end
 
       context "The job is started" do
@@ -58,14 +61,17 @@
 
         it 'should no longer show as waiting' do
           expect(job_state_tracker).not_to be_waiting(job_id)
+          expect(job_state_tracker.waiting_size).to be(0)
         end
 
         it 'should now show as running' do
           expect(job_state_tracker).to be_running(job_id)
+          expect(job_state_tracker.running_size).to be(1)
         end
 
         it 'should not show as finished' do
           expect(job_state_tracker).not_to be_finished(job_id)
+          expect(job_state_tracker.results_size).to be(0)
         end
 
         context "The job completes successfully" do
@@ -75,14 +81,17 @@
 
           it 'should not show as waiting' do
             expect(job_state_tracker).not_to be_waiting(job_id)
+            expect(job_state_tracker.waiting_size).to be(0)
           end
 
           it 'should no longer show as running' do
             expect(job_state_tracker).not_to be_running(job_id)
+            expect(job_state_tracker.running_size).to be(0)
           end
 
           it 'should show as finished' do
             expect(job_state_tracker).to be_finished(job_id)
+            expect(job_state_tracker.results_size).to be(1)
           end
 
           it 'should have a retrievable result' do
@@ -96,14 +105,17 @@
 
             it 'should not show as waiting' do
               expect(job_state_tracker).not_to be_waiting(job_id)
+              expect(job_state_tracker.waiting_size).to be(0)
             end
 
             it 'should not show as running' do
               expect(job_state_tracker).not_to be_running(job_id)
+              expect(job_state_tracker.running_size).to be(0)
             end
 
             it 'should no longer show as finished' do
               expect(job_state_tracker).not_to be_finished(job_id)
+              expect(job_state_tracker.results_size).to be(0)
             end
           end
         end
@@ -115,14 +127,17 @@
 
           it 'should not show as waiting' do
             expect(job_state_tracker).not_to be_waiting(job_id)
+            expect(job_state_tracker.waiting_size).to be(0)
           end
 
           it 'should no longer show as running' do
             expect(job_state_tracker).not_to be_running(job_id)
+            expect(job_state_tracker.running_size).to be(0)
           end
 
           it 'should show as finished' do
             expect(job_state_tracker).to be_finished(job_id)
+            expect(job_state_tracker.results_size).to be(1)
           end
 
           it 'should have a retrievable result' do
@@ -136,14 +151,17 @@
 
             it 'should not show as waiting' do
               expect(job_state_tracker).not_to be_waiting(job_id)
+              expect(job_state_tracker.waiting_size).to be(0)
             end
 
             it 'should not show as running' do
               expect(job_state_tracker).not_to be_running(job_id)
+              expect(job_state_tracker.running_size).to be(0)
             end
 
             it 'should no longer show as finished' do
               expect(job_state_tracker).not_to be_finished(job_id)
+              expect(job_state_tracker.results_size).to be(0)
             end
           end
         end