shoryuken 1.0.0 → 2.0.0

data/lib/shoryuken/message.rb

@@ -0,0 +1,70 @@
+module Shoryuken
+  class Message
+    attr_accessor :client, :queue_url, :queue_name, :data
+
+    def initialize(client, queue, data)
+      self.client = client
+      self.data = data
+
+      if queue.is_a?(Shoryuken::Queue)
+        self.queue_url = queue.url
+        self.queue_name = queue.name
+      else
+        # TODO: Remove next major release
+        Shoryuken.loggger.warn do
+          '[DEPRECATION] Passing a queue url into Shoryuken::Message is deprecated, please pass the queue itself'
+        end
+        self.queue_url = queue
+      end
+    end
+
+    def delete
+      client.delete_message(
+        queue_url: queue_url,
+        receipt_handle: data.receipt_handle
+      )
+    end
+
+    def change_visibility(options)
+      client.change_message_visibility(
+        options.merge(queue_url: queue_url, receipt_handle: data.receipt_handle)
+      )
+    end
+
+    def visibility_timeout=(timeout)
+      client.change_message_visibility(
+        queue_url: queue_url,
+        receipt_handle: data.receipt_handle,
+        visibility_timeout: timeout
+      )
+    end
+
+    def message_id
+      data.message_id
+    end
+
+    def receipt_handle
+      data.receipt_handle
+    end
+
+    def md5_of_body
+      data.md5_of_body
+    end
+
+    def body
+      data.body
+    end
+
+    def attributes
+      data.attributes
+    end
+
+    def md5_of_message_attributes
+      data.md5_of_message_attributes
+    end
+
+    def message_attributes
+      data.message_attributes
+    end
+  end
+end

data/lib/shoryuken/middleware/chain.rb

@@ -72,7 +72,7 @@ module Shoryuken
         i = entries.index { |entry| entry.klass == newklass }
         new_entry = i.nil? ? Entry.new(newklass, *args) : entries.delete_at(i)
         i = entries.find_index { |entry| entry.klass == oldklass } || entries.count - 1
-        entries.insert(i+1, new_entry)
+        entries.insert(i + 1, new_entry)
       end
 
       def exists?(klass)
@@ -117,7 +117,7 @@ module Shoryuken
 
       def patch_deprecated_middleware!(klass)
         if klass.instance_method(:call).arity == 3
-          Shoryuken.logger.warn "[DEPRECATION] #{klass.name}#call(worker_instance, queue, sqs_msg) is deprecated. Please use #{klass.name}#call(worker_instance, queue, sqs_msg, body)"
+          Shoryuken.logger.warn { "[DEPRECATION] #{klass.name}#call(worker_instance, queue, sqs_msg) is deprecated. Please use #{klass.name}#call(worker_instance, queue, sqs_msg, body)" }
 
           klass.class_eval do
             alias_method :deprecated_call, :call

data/lib/shoryuken/middleware/server/exponential_backoff_retry.rb

@@ -0,0 +1,50 @@
+module Shoryuken
+  module Middleware
+    module Server
+      class ExponentialBackoffRetry
+        include Util
+
+        def call(worker, queue, sqs_msg, body)
+          started_at = Time.now
+          yield
+        rescue
+          retry_intervals = Array(worker.class.get_shoryuken_options['retry_intervals'])
+
+          if retry_intervals.empty? || !handle_failure(sqs_msg, started_at, retry_intervals)
+            # Re-raise the exception if the job is not going to be exponential backoff retried.
+            # This allows custom middleware (like exception notifiers) to be aware of the unhandled failure.
+            raise
+          end
+        end
+
+        private
+
+        def handle_failure(sqs_msg, started_at, retry_intervals)
+          attempts = sqs_msg.attributes['ApproximateReceiveCount']
+
+          return unless attempts
+
+          attempts = attempts.to_i - 1
+
+          interval = if attempts < retry_intervals.size
+                       retry_intervals[attempts]
+                     else
+                       retry_intervals.last
+                     end
+
+          # Visibility timeouts are limited to a total 12 hours, starting from the receipt of the message.
+          # We calculate the maximum timeout by subtracting the amount of time since the receipt of the message.
+          #
+          # From the docs:  "Amazon SQS restarts the timeout period using the new value."
+          # http://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/AboutVT.html#AboutVT-extending-message-visibility-timeout
+          max_timeout = 43200 - (Time.now - started_at).ceil - 1
+          interval = max_timeout if interval > max_timeout
+
+          sqs_msg.change_visibility(visibility_timeout: interval.to_i)
+
+          logger.info { "Message #{sqs_msg.message_id} failed, will be retried in #{interval} seconds." }
+        end
+      end
+    end
+  end
+end

data/lib/shoryuken/middleware/server/timing.rb

@@ -9,19 +9,19 @@ module Shoryuken
             begin
               started_at = Time.now
 
-              logger.info "started at #{started_at}"
+              logger.info { "started at #{started_at}" }
 
               yield
 
               total_time = elapsed(started_at)
 
               if (total_time / 1000.0) > (timeout = Shoryuken::Client.queues(queue).visibility_timeout)
-                logger.warn "exceeded the queue visibility timeout by #{total_time - (timeout * 1000)} ms"
+                logger.warn { "exceeded the queue visibility timeout by #{total_time - (timeout * 1000)} ms" }
               end
 
-              logger.info "completed in: #{total_time} ms"
+              logger.info { "completed in: #{total_time} ms" }
             rescue => e
-              logger.info "failed in: #{elapsed(started_at)} ms"
+              logger.info { "failed in: #{elapsed(started_at)} ms" }
               raise e
             end
           end

data/lib/shoryuken/processor.rb

@@ -19,12 +19,10 @@ module Shoryuken
       timer = auto_visibility_timeout(queue, sqs_msg, worker.class)
 
       begin
-        defer do
-          body = get_body(worker.class, sqs_msg)
+        body = get_body(worker.class, sqs_msg)
 
-          worker.class.server_middleware.invoke(worker, queue, sqs_msg, body) do
-            worker.perform(sqs_msg, body)
-          end
+        worker.class.server_middleware.invoke(worker, queue, sqs_msg, body) do
+          worker.perform(sqs_msg, body)
         end
 
         @manager.async.processor_done(queue, current_actor)
@@ -35,22 +33,31 @@ module Shoryuken
 
     private
 
-    def auto_visibility_timeout(queue, sqs_msg, worker_class)
-      if worker_class.auto_visibility_timeout?
+    class MessageVisibilityExtender
+      include Celluloid
+      include Util
+
+      def auto_extend(queue, sqs_msg, worker_class)
         queue_visibility_timeout = Shoryuken::Client.queues(queue).visibility_timeout
 
-        timer = every(queue_visibility_timeout - 5) do
+        every(queue_visibility_timeout - 5) do
           begin
-            logger.debug "Extending message #{worker_name(worker_class, sqs_msg)}/#{queue}/#{sqs_msg.message_id} visibility timeout by #{queue_visibility_timeout}s."
+            logger.debug { "Extending message #{worker_name(worker_class, sqs_msg)}/#{queue}/#{sqs_msg.message_id} visibility timeout by #{queue_visibility_timeout}s." }
 
             sqs_msg.visibility_timeout = queue_visibility_timeout
           rescue => e
-            logger.error "Could not auto extend the message #{worker_class}/#{queue}/#{sqs_msg.message_id} visibility timeout. Error: #{e.message}"
+            logger.error { "Could not auto extend the message #{worker_class}/#{queue}/#{sqs_msg.message_id} visibility timeout. Error: #{e.message}" }
           end
         end
       end
+    end
 
-      timer
+    def auto_visibility_timeout(queue, sqs_msg, worker_class)
+      return unless worker_class.auto_visibility_timeout?
+
+      @visibility_extender ||= MessageVisibilityExtender.new_link
+
+      @visibility_extender.auto_extend(queue, sqs_msg, worker_class)
     end
 
     def get_body(worker_class, sqs_msg)
@@ -72,11 +79,18 @@ module Shoryuken
       when :text, nil
         sqs_msg.body
       else
-        body_parser.parse(sqs_msg.body) if body_parser.respond_to?(:parse) # i.e. JSON.parse(...)
+        if body_parser.respond_to?(:parse)
+          # JSON.parse
+          body_parser.parse(sqs_msg.body)
+        elsif body_parser.respond_to?(:load)
+          # see https://github.com/phstc/shoryuken/pull/91
+          # JSON.load
+          body_parser.load(sqs_msg.body)
+        end
       end
     rescue => e
-      logger.error "Error parsing the message body: #{e.message}\nbody_parser: #{body_parser}\nsqs_msg.body: #{sqs_msg.body}"
-      nil
+      logger.error { "Error parsing the message body: #{e.message}\nbody_parser: #{body_parser}\nsqs_msg.body: #{sqs_msg.body}" }
+      raise
     end
   end
 end

data/lib/shoryuken/queue.rb

@@ -0,0 +1,84 @@
+module Shoryuken
+  class Queue
+    attr_accessor :name, :client, :url
+
+    def initialize(client, name)
+      self.name   = name
+      self.client = client
+      self.url    = client.get_queue_url(queue_name: name).queue_url
+    end
+
+    def visibility_timeout
+      client.get_queue_attributes(
+        queue_url: url,
+        attribute_names: ['VisibilityTimeout']
+      ).attributes['VisibilityTimeout'].to_i
+    end
+
+    def delete_messages(options)
+      client.delete_message_batch(options.merge(queue_url: url))
+    end
+
+    def send_message(options)
+      options = sanitize_message!(options).merge(queue_url: url)
+
+      Shoryuken.client_middleware.invoke(options) do
+        client.send_message(options)
+      end
+    end
+
+    def send_messages(options)
+      client.send_message_batch(sanitize_messages!(options).merge(queue_url: url))
+    end
+
+    def receive_messages(options)
+      client.receive_message(options.merge(queue_url: url)).
+        messages.
+        map { |m| Message.new(client, self, m) }
+    end
+
+    private
+
+    def sanitize_messages!(options)
+      options = case
+                when options.is_a?(Array)
+                  { entries: options.map.with_index { |m, index| { id: index.to_s, message_body: m } } }
+                when options.is_a?(Hash)
+                  options
+                end
+
+      validate_messages!(options)
+
+      options
+    end
+
+    def sanitize_message!(options)
+      options = case
+                when options.is_a?(String)
+                  # send_message('message')
+                  { message_body: options }
+                when options.is_a?(Hash)
+                  options
+                end
+
+      validate_message!(options)
+
+      options
+    end
+
+    def validate_messages!(options)
+      options[:entries].map { |m| validate_message!(m) }
+    end
+
+    def validate_message!(options)
+      body = options[:message_body]
+      if body.is_a?(Hash)
+        options[:message_body] = JSON.dump(body)
+      elsif !body.is_a?(String)
+        fail ArgumentError, "The message body must be a String and you passed a #{body.class}"
+      end
+
+      options
+    end
+  end
+end

data/lib/shoryuken/sns_arn.rb

@@ -1,6 +1,6 @@
 module Shoryuken
   class SnsArn
-    def initialize topic
+    def initialize(topic)
       @topic = topic
     end
 
@@ -13,14 +13,14 @@ module Shoryuken
     def account_id
       Shoryuken::Client.account_id.tap do |account_id|
         if account_id.to_s.empty?
-          fail "To generate SNS ARNs, you must assign an :account_id in your Shoryuken::Client."
+          fail 'To generate SNS ARNs, you must assign an :account_id in your Shoryuken::Client.'
         end
       end
     end
 
     def region
       Aws.config.fetch(:region) do
-        fail "To generate SNS ARNs, you must include a :region in your AWS config."
+        fail 'To generate SNS ARNs, you must include a :region in your AWS config.'
       end
     end
   end

data/lib/shoryuken/util.rb

@@ -3,9 +3,9 @@ module Shoryuken
     def watchdog(last_words)
       yield
     rescue => ex
-      logger.error last_words
-      logger.error ex
-      logger.error ex.backtrace.join("\n")
+      logger.error { last_words }
+      logger.error { ex }
+      logger.error { ex.backtrace.join("\n") }
     end
 
     def logger
@@ -18,9 +18,8 @@ module Shoryuken
     end
 
     def unparse_queues(queues)
-      queues.inject({}) do |queue_and_weights, name|
+      queues.each_with_object({}) do |name, queue_and_weights|
         queue_and_weights[name] = queue_and_weights[name].to_i + 1
-        queue_and_weights
       end.to_a
     end
 
@@ -31,7 +30,7 @@ module Shoryuken
           && sqs_msg.message_attributes['shoryuken_class'] \
           && sqs_msg.message_attributes['shoryuken_class'][:string_value] == ActiveJob::QueueAdapters::ShoryukenAdapter::JobWrapper.to_s
 
-        "ActiveJob/#{body['job_class'].constantize}"
+        "ActiveJob/#{body['job_class']}"
       else
         worker_class.to_s
       end

data/lib/shoryuken/version.rb

@@ -1,3 +1,3 @@
 module Shoryuken
-  VERSION = '1.0.0'
+  VERSION = '2.0.0'
 end

data/lib/shoryuken/worker.rb

@@ -13,7 +13,6 @@ module Shoryuken
           data_type: 'String'
         }
 
-        body = JSON.dump(body) if body.is_a?(Hash)
         options[:message_body] = body
 
         Shoryuken::Client.queues(get_shoryuken_options['queue']).send_message(options)
@@ -40,7 +39,7 @@ module Shoryuken
       end
 
       def shoryuken_options(opts = {})
-        @shoryuken_options = get_shoryuken_options.merge(stringify_keys(Hash(opts)))
+        @shoryuken_options = get_shoryuken_options.merge(stringify_keys(opts || {}))
         queue = @shoryuken_options['queue']
         if queue.respond_to? :call
           queue = queue.call

data/lib/shoryuken.rb

@@ -1,6 +1,5 @@
 require 'yaml'
 require 'aws-sdk-core'
-require 'aws-sdk-resources'
 require 'time'
 
 require 'shoryuken/version'
@@ -8,12 +7,15 @@ require 'shoryuken/core_ext'
 require 'shoryuken/util'
 require 'shoryuken/logging'
 require 'shoryuken/environment_loader'
+require 'shoryuken/queue'
+require 'shoryuken/message'
 require 'shoryuken/client'
 require 'shoryuken/worker'
 require 'shoryuken/worker_registry'
 require 'shoryuken/default_worker_registry'
 require 'shoryuken/middleware/chain'
 require 'shoryuken/middleware/server/auto_delete'
+require 'shoryuken/middleware/server/exponential_backoff_retry'
 require 'shoryuken/middleware/server/timing'
 require 'shoryuken/sns_arn'
 require 'shoryuken/topic'
@@ -27,8 +29,9 @@ module Shoryuken
     timeout: 8
   }
 
-  @@queues          = []
+  @@queues = []
   @@worker_registry = DefaultWorkerRegistry.new
+  @@active_job_queue_name_prefixing = false
 
   class << self
     def options
@@ -55,11 +58,14 @@ module Shoryuken
       @@worker_registry
     end
 
-    # Shoryuken.configure_server do |config|
-    #   config.server_middleware do |chain|
-    #     chain.add MyServerHook
-    #   end
-    # end
+    def active_job_queue_name_prefixing
+      @@active_job_queue_name_prefixing
+    end
+
+    def active_job_queue_name_prefixing=(prefixing)
+      @@active_job_queue_name_prefixing = prefixing
+    end
+
     def configure_server
       yield self if server?
     end
@@ -70,12 +76,23 @@ module Shoryuken
       @server_chain
     end
 
+    def configure_client
+      yield self
+    end
+
+    def client_middleware
+      @client_chain ||= default_client_middleware
+      yield @client_chain if block_given?
+      @client_chain
+    end
+
     def default_worker_options
       @@default_worker_options ||= {
         'queue'                   => 'default',
         'delete'                  => false,
         'auto_delete'             => false,
         'auto_visibility_timeout' => false,
+        'retry_intervals'         => nil,
         'batch'                   => false }
     end
 
@@ -104,6 +121,7 @@ module Shoryuken
     def default_server_middleware
       Middleware::Chain.new do |m|
         m.add Middleware::Server::Timing
+        m.add Middleware::Server::ExponentialBackoffRetry
         m.add Middleware::Server::AutoDelete
         if defined?(::ActiveRecord::Base)
           require 'shoryuken/middleware/server/active_record'
@@ -112,6 +130,10 @@ module Shoryuken
       end
     end
 
+    def default_client_middleware
+      Middleware::Chain.new
+    end
+
     def server?
       defined?(Shoryuken::CLI)
     end

data/shoryuken.gemspec

@@ -4,27 +4,26 @@ $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
 require 'shoryuken/version'
 
 Gem::Specification.new do |spec|
-  spec.name          = "shoryuken"
+  spec.name          = 'shoryuken'
   spec.version       = Shoryuken::VERSION
-  spec.authors       = ["Pablo Cantero"]
-  spec.email         = ["pablo@pablocantero.com"]
-  spec.description = spec.summary = %q{Shoryuken is a super efficient AWS SQS thread based message processor}
-  spec.homepage      = "https://github.com/phstc/shoryuken"
-  spec.license       = "LGPL-3.0"
+  spec.authors       = ['Pablo Cantero']
+  spec.email         = ['pablo@pablocantero.com']
+  spec.description = spec.summary = %q(Shoryuken is a super efficient AWS SQS thread based message processor)
+  spec.homepage      = 'https://github.com/phstc/shoryuken'
+  spec.license       = 'LGPL-3.0'
 
   spec.files         = `git ls-files -z`.split("\x0")
   spec.executables   = %w[shoryuken]
   spec.test_files    = spec.files.grep(%r{^(test|spec|features)/})
-  spec.require_paths = ["lib"]
+  spec.require_paths = ['lib']
 
-  spec.add_development_dependency "bundler", "~> 1.6"
-  spec.add_development_dependency "rake"
-  spec.add_development_dependency "rspec"
-  spec.add_development_dependency "pry-byebug"
-  spec.add_development_dependency "nokogiri"
-  spec.add_development_dependency "dotenv"
+  spec.add_development_dependency 'bundler', '~> 1.6'
+  spec.add_development_dependency 'rake'
+  spec.add_development_dependency 'rspec'
+  spec.add_development_dependency 'pry-byebug'
+  spec.add_development_dependency 'nokogiri'
+  spec.add_development_dependency 'dotenv'
 
-  spec.add_dependency "aws-sdk-core", "2.0.21"
-  spec.add_dependency "aws-sdk-resources", "2.0.21.pre"
-  spec.add_dependency "celluloid", "~> 0.16.0"
+  spec.add_dependency 'aws-sdk-core', '~> 2.0.21'
+  spec.add_dependency 'celluloid', '~> 0.16.0'
 end

data/spec/shoryuken/client_spec.rb

@@ -5,12 +5,13 @@ describe Shoryuken::Client do
   let(:sqs)         { Aws::SQS::Client.new(stub_responses: true, credentials: credentials) }
   let(:queue_name)  { 'shoryuken' }
   let(:queue_url)   { 'https://eu-west-1.amazonaws.com:6059/123456789012/shoryuken' }
-
-  before do
-    described_class.sqs = sqs
-  end
+  let(:sqs_endpoint) { 'http://localhost:4568' }
+  let(:sns_endpoint) { 'http://0.0.0.0:4568'   }
 
   describe '.queue' do
+    before do
+      described_class.sqs = sqs
+    end
     it 'memoizes queues' do
       sqs.stub_responses(:get_queue_url, { queue_url: queue_url }, { queue_url: 'xyz' })
 
@@ -18,4 +19,44 @@ describe Shoryuken::Client do
       expect(Shoryuken::Client.queues(queue_name).url).to eq queue_url
     end
   end
+
+  describe 'environment variable endpoints' do
+    before do
+      ENV['AWS_SQS_ENDPOINT'] = sqs_endpoint
+      ENV['AWS_SNS_ENDPOINT'] = sns_endpoint
+      ENV['AWS_REGION'] = 'us-east-1'
+      Shoryuken.options[:aws] = {}
+    end
+
+    it 'will use config file settings if set' do
+      load_config_file_by_file_name('shoryuken_endpoint.yml')
+      expect(described_class.sqs.config.endpoint.to_s).to eql('https://github.com/phstc/shoryuken:4568')
+      expect(described_class.sns.config.endpoint.to_s).to eq('http://127.0.0.1:4568')
+    end
+
+    it 'should fallback to environment variable if config file not found or set' do
+      load_config_file_by_file_name(nil)
+      expect(described_class.sqs.config.endpoint.to_s).to eql(sqs_endpoint)
+      expect(described_class.sns.config.endpoint.to_s).to eq(sns_endpoint)
+    end
+
+    it 'should fallback to environment variable if config file found but settings not set' do
+      load_config_file_by_file_name('shoryuken.yml')
+      expect(described_class.sqs.config.endpoint.to_s).to eql(sqs_endpoint)
+      expect(described_class.sns.config.endpoint.to_s).to eq(sns_endpoint)
+    end
+
+    it 'will fallback to default settings if no config file settings or environment variables found' do
+      ENV['AWS_SQS_ENDPOINT'] = nil
+      ENV['AWS_SNS_ENDPOINT'] = nil
+      load_config_file_by_file_name('shoryuken.yml')
+      expect(described_class.sqs.config.endpoint.to_s).to eql('https://sqs.us-east-1.amazonaws.com')
+      expect(described_class.sns.config.endpoint.to_s).to eq('https://sns.us-east-1.amazonaws.com')
+    end
+  end
+
+  def load_config_file_by_file_name(file_name)
+    path_name = file_name ? File.join(File.expand_path('../../..', __FILE__), 'spec', file_name) : nil
+    Shoryuken::EnvironmentLoader.load(config_file: path_name)
+  end
 end

data/spec/shoryuken/default_worker_registry_spec.rb

@@ -48,7 +48,7 @@ describe Shoryuken::DefaultWorkerRegistry do
         string_value: explicit_worker.to_s,
         data_type: 'String' } if explicit_worker
 
-      double Aws::SQS::Message,
+      double Shoryuken::Message,
         body: 'test',
         message_attributes: attributes,
         message_id: SecureRandom.uuid

data/spec/shoryuken/fetcher_spec.rb

@@ -4,11 +4,11 @@ require 'shoryuken/fetcher'
 
 describe Shoryuken::Fetcher do
   let(:manager)    { double Shoryuken::Manager }
-  let(:queue)      { double Aws::SQS::Queue }
+  let(:queue)      { double Shoryuken::Queue }
   let(:queue_name) { 'default' }
 
   let(:sqs_msg) do
-    double Aws::SQS::Message,
+    double Shoryuken::Message,
       queue_url: queue_name,
       body: 'test',
       message_id: 'fc754df7-9cc2-4c41-96ca-5996a44b771e'
@@ -24,7 +24,7 @@ describe Shoryuken::Fetcher do
 
   describe '#fetch' do
     it 'calls pause when no message' do
-      allow(queue).to receive(:receive_messages).with(max_number_of_messages: 1, message_attribute_names: ['All']).and_return([])
+      allow(queue).to receive(:receive_messages).with(max_number_of_messages: 1, attribute_names: ['All'], message_attribute_names: ['All']).and_return([])
 
       expect(manager).to receive(:pause_queue!).with(queue_name)
       expect(manager).to receive(:dispatch)
@@ -33,7 +33,7 @@ describe Shoryuken::Fetcher do
     end
 
     it 'assigns messages' do
-      allow(queue).to receive(:receive_messages).with(max_number_of_messages: 5, message_attribute_names: ['All']).and_return(sqs_msg)
+      allow(queue).to receive(:receive_messages).with(max_number_of_messages: 5, attribute_names: ['All'], message_attribute_names: ['All']).and_return(sqs_msg)
 
       expect(manager).to receive(:rebalance_queue_weight!).with(queue_name)
       expect(manager).to receive(:assign).with(queue_name, sqs_msg)
@@ -45,7 +45,7 @@ describe Shoryuken::Fetcher do
     it 'assigns messages in batch' do
       TestWorker.get_shoryuken_options['batch'] = true
 
-      allow(queue).to receive(:receive_messages).with(max_number_of_messages: described_class::FETCH_LIMIT, message_attribute_names: ['All']).and_return(sqs_msg)
+      allow(queue).to receive(:receive_messages).with(max_number_of_messages: described_class::FETCH_LIMIT, attribute_names: ['All'], message_attribute_names: ['All']).and_return(sqs_msg)
 
       expect(manager).to receive(:rebalance_queue_weight!).with(queue_name)
       expect(manager).to receive(:assign).with(queue_name, [sqs_msg])
@@ -58,7 +58,7 @@ describe Shoryuken::Fetcher do
       let(:queue_name) { 'notfound' }
 
       it 'ignores batch' do
-        allow(queue).to receive(:receive_messages).with(max_number_of_messages: 5, message_attribute_names: ['All']).and_return(sqs_msg)
+        allow(queue).to receive(:receive_messages).with(max_number_of_messages: 5, attribute_names: ['All'], message_attribute_names: ['All']).and_return(sqs_msg)
 
         expect(manager).to receive(:rebalance_queue_weight!).with(queue_name)
         expect(manager).to receive(:assign).with(queue_name, sqs_msg)