chore-core 1.10.0 → 4.0.0

data/lib/chore/unit_of_work.rb

@@ -2,13 +2,14 @@ module Chore
   # Simple class to hold job processing information.
   # Has six attributes:
   # * +:id+ The queue implementation specific identifier for this message.
+  # * +:receipt_handle+ The queue implementation specific identifier for the receipt of this message.
   # * +:queue_name+ The name of the queue the job came from
   # * +:queue_timeout+ The time (in seconds) before the job will get re-enqueued if not processed
   # * +:message+ The actual data of the message.
   # * +:previous_attempts+ The number of times the work has been attempted previously.
   # * +:consumer+ The consumer instance used to fetch this message. Most queue implementations won't need access to this, but some (RabbitMQ) will. So we
   # make sure to pass it along with each message. This instance will be used by the Worker for things like <tt>complete</tt> and </tt>reject</tt>.
-  class UnitOfWork < Struct.new(:id,:queue_name,:queue_timeout,:message,:previous_attempts,:consumer,:decoded_message, :klass)
+  class UnitOfWork < Struct.new(:id, :receipt_handle, :queue_name, :queue_timeout, :message, :previous_attempts, :consumer, :decoded_message, :klass)
     # The time at which this unit of work was created
     attr_accessor :created_at
 

data/lib/chore/util.rb

@@ -2,7 +2,7 @@ module Chore
 
   # Collection of utilities and helpers used by Chore internally
   module Util
-    
+
     # To avoid bringing in all of active_support, we implemented constantize here
     def constantize(camel_cased_word)
       names = camel_cased_word.split('::')
@@ -14,5 +14,9 @@ module Chore
       end
       constant
     end
+
+    def procline(str)
+      $0 = str
+    end
   end
 end

data/lib/chore/version.rb

@@ -1,7 +1,7 @@
 module Chore
   module Version #:nodoc:
-    MAJOR = 1
-    MINOR = 10
+    MAJOR = 4
+    MINOR = 0
     PATCH = 0
 
     STRING = [ MAJOR, MINOR, PATCH ].join('.')

data/lib/chore/worker.rb

@@ -42,6 +42,28 @@ module Chore
       @started_at + total_timeout
     end
 
+    def duplicate_work?(item)
+      # if we've got a duplicate, remove the message from the queue by not actually running and also not reporting any errors
+      payload = options[:payload_handler].payload(item.decoded_message)
+
+       # if we're hitting the custom dedupe key, we want to remove this message from the queue
+      if item.klass.has_dedupe_lambda?
+        dedupe_key = item.klass.dedupe_key(*payload)
+        if dedupe_key.nil? || dedupe_key.strip.empty? # if the dedupe key is nil, don't continue with the rest of the dedupe lambda logic
+          Chore.logger.info { "#{item.klass} dedupe key nil, skipping memcached lookup." }
+          return false
+        end
+
+        if item.consumer.duplicate_message?(dedupe_key, item.klass, item.queue_timeout)
+          Chore.logger.info { "Found and deleted duplicate job #{item.klass}"}
+          item.consumer.complete(item.id, item.receipt_handle)
+          return true
+        end
+      end
+
+      return false
+    end
+
     # The workhorse. Do the work, all of it. This will block for an entirely unspecified amount
     # of time based on the work to be performed. This will:
     # * Decode each message.
@@ -58,12 +80,16 @@ module Chore
         begin
           item.decoded_message = options[:payload_handler].decode(item.message)
           item.klass = options[:payload_handler].payload_class(item.decoded_message)
+
+          next if duplicate_work?(item)
+
+          Chore.run_hooks_for(:worker_to_start, item)
           start_item(item)
         rescue => e
           Chore.logger.error { "Failed to run job for #{item.message} with error: #{e.message} #{e.backtrace * "\n"}" }
           if item.current_attempt >= Chore.config.max_attempts
             Chore.run_hooks_for(:on_permanent_failure,item.queue_name,item.message,e)
-            item.consumer.complete(item.id)
+            item.consumer.complete(item.id, item.receipt_handle)
           else
             Chore.run_hooks_for(:on_failure,item.message,e)
             item.consumer.reject(item.id)
@@ -86,9 +112,10 @@ module Chore
       begin
         Chore.logger.info { "Running job #{klass} with params #{message}"}
         perform_job(klass,message)
-        item.consumer.complete(item.id)
+        item.consumer.complete(item.id, item.receipt_handle)
         Chore.logger.info { "Finished job #{klass} with params #{message}"}
         klass.run_hooks_for(:after_perform, message)
+        Chore.run_hooks_for(:worker_ended, item)
       rescue Job::RejectMessageException
         item.consumer.reject(item.id)
         Chore.logger.error { "Failed to run job for #{item.message}  with error: Job raised a RejectMessageException" }
@@ -114,7 +141,7 @@ module Chore
       Chore.logger.error { "Failed to run job #{item.message} with error: #{e.message} at #{e.backtrace * "\n"}" }
       if item.current_attempt >= klass.options[:max_attempts]
         klass.run_hooks_for(:on_permanent_failure,item.queue_name,message,e)
-        item.consumer.complete(item.id)
+        item.consumer.complete(item.id, item.receipt_handle)
       else
         klass.run_hooks_for(:on_failure, message, e)
         item.consumer.reject(item.id)

data/spec/chore/cli_spec.rb

@@ -201,8 +201,8 @@ describe Chore::CLI do
 
       context 'given no value' do
         let(:command) { [] }
-        it 'is the default value, nil' do
-          subject.should == nil
+        it 'is the default value, 1' do
+          subject.should == 1
         end
       end
     end

data/spec/chore/consumer_spec.rb

@@ -22,10 +22,6 @@ describe Chore::Consumer do
     Chore::Consumer.should respond_to :reset_connection!
   end
 
-  it 'should have a class level cleanup method' do
-    Chore::Consumer.should respond_to :cleanup
-  end
-
   it 'should not have an implemented consume method' do
     expect { consumer.consume }.to raise_error(NotImplementedError)
   end
@@ -35,6 +31,6 @@ describe Chore::Consumer do
   end
 
   it 'should not have an implemented complete method' do
-    expect { consumer.complete(message) }.to raise_error(NotImplementedError)
+    expect { consumer.complete(message, nil) }.to raise_error(NotImplementedError)
   end
 end

data/spec/chore/duplicate_detector_spec.rb

@@ -2,7 +2,21 @@ require 'spec_helper'
 require 'securerandom'
 
 describe Chore::DuplicateDetector do
-  let(:memcache) { double("memcache") }
+  class FakeDalli
+    def initialize
+      @store = {}
+    end
+    def add(id, val, ttl=0)
+      if @store[id] && @store[id][:inserted] + @store[id][:ttl] > Time.now.to_i
+        return false
+      else
+        @store[id] = {:val => val, :ttl => ttl, :inserted => Time.now.to_i}
+        return true
+      end
+    end
+  end
+
+  let(:memcache) { FakeDalli.new }
   let(:dupe_on_cache_failure) { false }
   let(:dedupe_params)  { { :memcached_client => memcache, :dupe_on_cache_failure => dupe_on_cache_failure } }
   let(:dedupe) { Chore::DuplicateDetector.new(dedupe_params)}
@@ -15,12 +29,11 @@ describe Chore::DuplicateDetector do
 
   describe "#found_duplicate" do
     it 'should not return true if the message has not already been seen' do
-      expect(memcache).to receive(:add).and_return(true)
       expect(dedupe.found_duplicate?(message_data)).to_not be true
     end
 
     it 'should return true if the message has already been seen' do
-      expect(memcache).to receive(:add).and_return(false)
+      memcache.add(message_data[:id], 1, message_data[:visibility_timeout])
       expect(dedupe.found_duplicate?(message_data)).to be true
     end
 
@@ -34,13 +47,12 @@ describe Chore::DuplicateDetector do
     end
 
     it "should set the timeout to be the queue's " do
-      expect(memcache).to receive(:add).with(id,"1",timeout).and_return(true)
+      expect(memcache).to receive(:add).with(id,"1",timeout).and_call_original
       expect(dedupe.found_duplicate?(message_data)).to be false
     end
 
     it "should call #visibility_timeout once and only once" do
       expect(queue).to receive(:visibility_timeout).once
-      expect(memcache).to receive(:add).at_least(3).times.and_return(true)
       3.times { dedupe.found_duplicate?(message_data) }
     end
 

data/spec/chore/fetcher_spec.rb

@@ -35,15 +35,4 @@ describe Chore::Fetcher do
       fetcher.start
     end
   end
-
-  describe "cleaning up" do
-    before(:each) do
-      manager.stub(:assign)
-    end
-    
-    it "should run cleanup on each queue" do
-      consumer.should_receive(:cleanup).with('test')
-      fetcher.start
-    end
-  end
 end

data/spec/chore/manager_spec.rb

@@ -35,6 +35,13 @@ describe Chore::Manager do
         manager.assign(work)
       end
     end
+
+    describe 'returning work' do
+      it 'should return work to the fetcher' do
+        expect(fetcher).to receive(:return_work).with([work])
+        manager.return_work([work])
+      end
+    end
   end
 
 end

data/spec/chore/queues/filesystem/filesystem_consumer_spec.rb

@@ -8,10 +8,16 @@ describe Chore::Queues::Filesystem::Consumer do
   let(:publisher) { Chore::Queues::Filesystem::Publisher.new }
   let(:test_queues_dir) { "test-queues" }
   let(:test_queue) { "test-queue" }
+  let(:default_timeout) { 60 }
+  let(:timeout) { nil }
 
   before do
     Chore.config.fs_queue_root = test_queues_dir
-    expect(Chore.config).to receive(:default_queue_timeout).and_return(60)
+    if timeout
+      File.open("#{config_dir}/timeout", "w") {|f| f << timeout.to_s}
+    else
+      expect(Chore.config).to receive(:default_queue_timeout).and_return(default_timeout)
+    end
     allow(consumer).to receive(:sleep)
   end
 
@@ -22,35 +28,79 @@ describe Chore::Queues::Filesystem::Consumer do
   let(:test_job_hash) {{:class => "TestClass", :args => "test-args"}}
   let(:new_dir) { described_class.new_dir(test_queue) }
   let(:in_progress_dir) { described_class.in_progress_dir(test_queue) }
+  let(:config_dir) { described_class.config_dir(test_queue) }
 
   describe ".cleanup" do
-    it "should move in_progress jobs to new dir" do
+    it "should move expired in_progress jobs to new dir" do
+      timestamp = Time.now.to_i - 1
+
+      FileUtils.touch("#{in_progress_dir}/foo.1.#{timestamp}.job")
+      described_class.cleanup(Time.now.to_i, new_dir, in_progress_dir)
+      expect(File.exist?("#{new_dir}/foo.2.job")).to eq(true)
+    end
+
+    it "should move non-timestamped jobs from in_progress_dir to new dir" do
       FileUtils.touch("#{in_progress_dir}/foo.1.job")
-      described_class.cleanup(test_queue)
+      described_class.cleanup(Time.now.to_i, new_dir, in_progress_dir)
       expect(File.exist?("#{new_dir}/foo.2.job")).to eq(true)
     end
+
+    it "should not affect non-expired jobs" do
+      timestamp = Time.now.to_i - 1
+
+      FileUtils.touch("#{in_progress_dir}/foo.1.#{timestamp}.job")
+      described_class.cleanup(Time.now.to_i - 2, new_dir, in_progress_dir)
+      expect(File.exist?("#{new_dir}/foo.2.job")).to eq(false)
+    end
   end
 
   describe ".make_in_progress" do
-    it "should move job to in_progress dir" do
+    it "should move non-empty job to in_progress dir" do
+      now = Time.now
+
+      Timecop.freeze(now) do
+        File.open("#{new_dir}/foo.1.job", "w") {|f| f << "{}"}
+        described_class.make_in_progress("foo.1.job", new_dir, in_progress_dir, default_timeout)
+        expect(File.exist?("#{in_progress_dir}/foo.1.#{now.to_i}.job")).to eq(true)
+      end
+    end
+
+    it "should not move empty jobs to in_progress dir" do
+      now = Time.now
+
+      Timecop.freeze(now) do
+        FileUtils.touch("#{new_dir}/foo.1.job")
+        described_class.make_in_progress("foo.1.job", new_dir, in_progress_dir, default_timeout)
+        expect(File.exist?("#{new_dir}/foo.1.job")).to eq(true)
+        expect(File.exist?("#{in_progress_dir}/foo.1.#{now.to_i}.job")).to eq(false)
+      end
+    end
+
+    it "should delete expired empty jobs" do
       FileUtils.touch("#{new_dir}/foo.1.job")
-      described_class.make_in_progress("foo.1.job", new_dir, in_progress_dir)
-      expect(File.exist?("#{in_progress_dir}/foo.1.job")).to eq(true)
+
+      now = Time.now + default_timeout
+      Timecop.freeze(now) do
+        described_class.make_in_progress("foo.1.job", new_dir, in_progress_dir, default_timeout)
+        expect(File.exist?("#{new_dir}/foo.1.job")).to eq(false)
+        expect(File.exist?("#{in_progress_dir}/foo.1.#{now.to_i}.job")).to eq(false)
+      end
     end
   end
 
   describe ".make_new_again" do
     it "should move job to new dir" do
-      FileUtils.touch("#{in_progress_dir}/foo.1.job")
-      described_class.make_new_again("foo.1.job", new_dir, in_progress_dir)
+      timestamp = Time.now.to_i
+      FileUtils.touch("#{in_progress_dir}/foo.1.#{timestamp}.job")
+      described_class.make_new_again("foo.1.#{timestamp}.job", new_dir, in_progress_dir)
       expect(File.exist?("#{new_dir}/foo.2.job")).to eq(true)
     end
   end
 
-  describe ".job_files" do
+  describe ".each_file" do
     it "should list jobs in dir" do
       FileUtils.touch("#{new_dir}/foo.1.job")
-      expect(described_class.job_files(new_dir)).to eq(["foo.1.job"])
+      expect {|b| described_class.each_file(new_dir, &b) }.to yield_with_args("foo.1.job")
     end
   end
 
@@ -71,7 +121,7 @@ describe Chore::Queues::Filesystem::Consumer do
       end
 
       it "should consume a published job and yield the job to the handler block" do
-        expect { |b| consumer.consume(&b) }.to yield_with_args(anything, 'test-queue', 60, test_job_hash.to_json, 0)
+        expect { |b| consumer.consume(&b) }.to yield_with_args(anything, anything, 'test-queue', 60, test_job_hash.to_json, 0)
       end
 
       context "rejecting a job" do
@@ -85,7 +135,9 @@ describe Chore::Queues::Filesystem::Consumer do
           end
           expect(rejected).to be true
 
-          expect { |b| consumer.consume(&b) }.to yield_with_args(anything, 'test-queue', 60, test_job_hash.to_json, 1)
+          Timecop.freeze(Time.now + 61) do
+            expect { |b| consumer.consume(&b) }.to yield_with_args(anything, anything, 'test-queue', 60, test_job_hash.to_json, 1)
+          end
         end
       end
 
@@ -93,16 +145,23 @@ describe Chore::Queues::Filesystem::Consumer do
         let!(:consumer_run_for_two_messages) { allow(consumer).to receive(:running?).and_return(true, false,true,false) }
 
         it "should remove job on completion" do
-          completed = false
+
           consumer.consume do |job_id, queue_name, job_hash|
+            expect(File).to receive(:delete).with(kind_of(String))
             consumer.complete(job_id)
-            completed = true
           end
-          expect(completed).to be true
 
           expect { |b| consumer.consume(&b) }.to_not yield_control
         end
       end
+
+      context "with queue-specific timeout config" do
+        let(:timeout) { 30 }
+
+        it "should consume a published job and yield the job to the handler block" do
+          expect { |b| consumer.consume(&b) }.to yield_with_args(anything, anything, 'test-queue', 30, test_job_hash.to_json, 0)
+        end
+      end
     end
 
     context "not finding a published job" do
@@ -112,4 +171,3 @@ describe Chore::Queues::Filesystem::Consumer do
     end
   end
 end
-

data/spec/chore/queues/sqs/consumer_spec.rb

@@ -1,133 +1,170 @@
 require 'spec_helper'
 
 describe Chore::Queues::SQS::Consumer do
-  let(:queue_name) { "test" }
-  let(:queue_url) { "test_url" }
-  let(:queues) { double("queues") }
-  let(:queue) { double("test_queue", :visibility_timeout=>10, :url=>"test_queue", :name=>"test_queue") }
+  include_context 'fake objects'
+
   let(:options) { {} }
   let(:consumer) { Chore::Queues::SQS::Consumer.new(queue_name) }
-  let(:message) { TestMessage.new("handle",queue, "message body", 1) }
-  let(:message_data) {{:id=>message.id, :queue=>message.queue.url, :visibility_timeout=>message.queue.visibility_timeout}}
-  let(:pool) { double("pool") }
-  let(:sqs) { double('AWS::SQS') }
-  let(:backoff_func) { nil }
+  let(:job) { {'class' => 'TestJob', 'args'=>[1,2,'3']} }
+  let(:backoff_func) { Proc.new { 2 + 2 } }
+
+  let(:receive_message_result) { Aws::SQS::Message::Collection.new([message], size: 1) }
+
+  let(:message) do
+      Aws::SQS::Message.new(
+      message_id: 'message id',
+      receipt_handle: "receipt_handle",
+      body: job.to_json,
+      data: job,
+      queue: queue,
+      queue_url: queue_url,
+    )
+  end
 
-  before do
-    allow(AWS::SQS).to receive(:new).and_return(sqs)
-    allow(sqs).to receive(:queues) { queues }
+  # Since a message handler is required (but not validated), this convenience method lets us
+  # effectively stub the block.
+  def consume(&block)
+    block = Proc.new{} unless block_given?
+    consumer.consume(&block)
+  end
 
-    allow(queues).to receive(:url_for) { queue_url }
-    allow(queues).to receive(:[]) { queue }
-    allow(queue).to receive(:receive_message) { message }
-    allow(pool).to receive(:empty!) { nil }
+  before do
+    allow(Aws::SQS::Client).to receive(:new).and_return(sqs)
+    allow(Aws::SQS::Queue).to receive(:new).and_return(queue)
+    allow(queue).to receive(:receive_messages).and_return(receive_message_result)
+    allow(message).to receive(:attributes).and_return({ 'ApproximateReceiveCount' => rand(10) })
   end
 
   describe "consuming messages" do
-    let!(:consumer_run_for_one_message) { allow(consumer).to receive(:running?).and_return(true, false) }
-    let!(:messages_be_unique) { allow_any_instance_of(Chore::DuplicateDetector).to receive(:found_duplicate?).and_return(false) }
-    let!(:queue_contain_messages) { allow(queue).to receive(:receive_messages).and_return(message) }
-
-    it 'should configure sqs' do
-      allow(Chore.config).to receive(:aws_access_key).and_return('key')
-      allow(Chore.config).to receive(:aws_secret_key).and_return('secret')
-
-      expect(AWS::SQS).to receive(:new).with(
-        :access_key_id => 'key',
-        :secret_access_key => 'secret',
-        :logger => Chore.logger,
-        :log_level => :debug
-      ).and_return(sqs)
-      consumer.consume
+    before do
+      allow(consumer).to receive(:running?).and_return(true, false)
     end
 
-    it 'should not configure sqs multiple times' do
-      allow(consumer).to receive(:running?).and_return(true, true, false)
+    context "should create objects for interacting with the SQS API" do
+      it 'should create an sqs client' do
+        expect(queue).to receive(:receive_messages)
+        consume
+      end
 
-      expect(AWS::SQS).to receive(:new).once.and_return(sqs)
-      consumer.consume
-    end
+      it "should only create an sqs client when one doesn't exist" do
+        allow(consumer).to receive(:running?).and_return(true, true, true, true, false, true, true)
+        expect(Aws::SQS::Client).to receive(:new).exactly(:once)
+        consume
+      end
 
-    it 'should look up the queue url based on the queue name' do
-      expect(queues).to receive(:url_for).with('test').and_return(queue_url)
-      consumer.consume
-    end
+      it 'should look up the queue url based on the queue name' do
+        expect(sqs).to receive(:get_queue_url).with(queue_name: queue_name)
+        consume
+      end
 
-    it 'should look up the queue based on the queue url' do
-      expect(queues).to receive(:[]).with(queue_url).and_return(queue)
-      consumer.consume
+      it 'should create a queue object' do
+        expect(consumer.send(:queue)).to_not be_nil
+        consume
+      end
     end
 
     context "should receive a message from the queue" do
-
       it 'should use the default size of 10 when no queue_polling_size is specified' do
-        expect(queue).to receive(:receive_messages).with(:limit => 10, :attributes => [:receive_count])
-        consumer.consume
+        expect(queue).to receive(:receive_messages).with(
+          :max_number_of_messages => 10,
+          :attribute_names => ['ApproximateReceiveCount']
+        ).and_return(message)
+        consume
       end
 
       it 'should respect the queue_polling_size when specified' do
         allow(Chore.config).to receive(:queue_polling_size).and_return(5)
-        expect(queue).to receive(:receive_messages).with(:limit => 5, :attributes => [:receive_count])
-        consumer.consume
+        expect(queue).to receive(:receive_messages).with(
+          :max_number_of_messages => 5,
+          :attribute_names => ['ApproximateReceiveCount']
+        )
+        consume
       end
     end
 
-    it "should check the uniqueness of the message" do
-      allow_any_instance_of(Chore::DuplicateDetector).to receive(:found_duplicate?).with(message_data).and_return(false)
-      consumer.consume
-    end
-
-    it "should yield the message to the handler block" do
-      expect { |b| consumer.consume(&b) }.to yield_with_args('handle', queue_name, 10, 'message body', 0)
-    end
-
-    it 'should not yield for a dupe message' do
-      allow_any_instance_of(Chore::DuplicateDetector).to receive(:found_duplicate?).with(message_data).and_return(true)
-      expect {|b| consumer.consume(&b) }.not_to yield_control
-    end
-
     context 'with no messages' do
-      let!(:consumer_run_for_one_message) { allow(consumer).to receive(:running?).and_return(true, true, false) }
-      let!(:queue_contain_messages) { allow(queue).to receive(:receive_messages).and_return(message, nil) }
+      before do
+        allow(consumer).to receive(:handle_messages).and_return([])
+      end
 
       it 'should sleep' do
         expect(consumer).to receive(:sleep).with(1)
-        consumer.consume
+        consume
       end
     end
 
     context 'with messages' do
-      let!(:consumer_run_for_one_message) { allow(consumer).to receive(:running?).and_return(true, true, false) }
-      let!(:queue_contain_messages) { allow(queue).to receive(:receive_messages).and_return(message, message) }
+      before do
+        allow(consumer).to receive(:duplicate_message?).and_return(false)
+        allow(queue).to receive(:receive_messages).and_return(message)
+      end
+
+      it "should check the uniqueness of the message" do
+        expect(consumer).to receive(:duplicate_message?)
+        consume
+      end
+
+      it "should yield the message to the handler block" do
+        expect { |b| consume(&b) }
+          .to yield_with_args(
+                message.message_id,
+                message.receipt_handle,
+                queue_name,
+                queue.attributes['VisibilityTimeout'].to_i,
+                message.body,
+                message.attributes['ApproximateReceiveCount'].to_i - 1
+              )
+      end
 
       it 'should not sleep' do
         expect(consumer).to_not receive(:sleep)
-        consumer.consume
+        consume
       end
+
+      context 'with duplicates' do
+        before do
+          allow(consumer).to receive(:duplicate_message?).and_return(true)
+        end
+
+        it 'should not yield for a dupe message' do
+          expect {|b| consume(&b) }.not_to yield_control
+        end
+      end
+    end
+  end
+
+  describe "completing work" do
+    it 'deletes the message from the queue' do
+      expect(queue).to receive(:delete_messages).with(entries: [{id: message.message_id, receipt_handle: message.receipt_handle}])
+      consumer.complete(message.message_id, message.receipt_handle)
     end
   end
 
   describe '#delay' do
-    let(:item) { Chore::UnitOfWork.new(message.id, message.queue, 60, message.body, 0, consumer) }
-    let(:backoff_func) { lambda { |item| 2 } }
+    let(:item) { Chore::UnitOfWork.new(message.message_id, message.receipt_handle, message.queue, 60, message.body, 0, consumer) }
+    let(:entries) do
+      [
+        { id: item.id, receipt_handle: item.receipt_handle, visibility_timeout: backoff_func.call(item) },
+      ]
+    end
 
     it 'changes the visiblity of the message' do
-      expect(queue).to receive(:batch_change_visibility).with(2, [item.id])
+      expect(queue).to receive(:change_message_visibility_batch).with(entries: entries)
       consumer.delay(item, backoff_func)
     end
   end
 
   describe '#reset_connection!' do
     it 'should reset the connection after a call to reset_connection!' do
-      expect(AWS::Core::Http::ConnectionPool).to receive(:pools).and_return([pool])
-      expect(pool).to receive(:empty!)
+      expect(Aws).to receive(:empty_connection_pools!)
       Chore::Queues::SQS::Consumer.reset_connection!
       consumer.send(:queue)
     end
 
     it 'should not reset the connection between calls' do
-      sqs = consumer.send(:queue)
-      expect(sqs).to be consumer.send(:queue)
+      expect(Aws).to receive(:empty_connection_pools!).once
+      q = consumer.send(:queue)
+      expect(consumer.send(:queue)).to be(q)
     end
 
     it 'should reconfigure sqs' do
@@ -135,13 +172,15 @@ describe Chore::Queues::SQS::Consumer do
       allow_any_instance_of(Chore::DuplicateDetector).to receive(:found_duplicate?).and_return(false)
 
       allow(queue).to receive(:receive_messages).and_return(message)
-      consumer.consume
+      allow(sqs).to receive(:receive_message).with({:attribute_names=>["ApproximateReceiveCount"], :max_number_of_messages=>10, :queue_url=>queue_url})
+
+      consume
 
       Chore::Queues::SQS::Consumer.reset_connection!
-      allow(AWS::SQS).to receive(:new).and_return(sqs)
+      allow(Aws::SQS::Client).to receive(:new).and_return(sqs)
 
       expect(consumer).to receive(:running?).and_return(true, false)
-      consumer.consume
+      consume
     end
   end
 end