Respect VisibilityTimeout when calling receive_message

Previously, receive_message would always use the queue's default
visibility timeout instead of the value passed as a query parameter when
calling the receive_message method on an SQS connection.

moto/sqs/models.py

@@ -265,7 +265,7 @@ class SQSBackend(BaseBackend):
 
         return message
 
-    def receive_messages(self, queue_name, count, wait_seconds_timeout):
+    def receive_messages(self, queue_name, count, wait_seconds_timeout, visibility_timeout):
         """
         Attempt to retrieve visible messages from a queue.
 
@@ -276,6 +276,7 @@ class SQSBackend(BaseBackend):
 
         :param string queue_name: The name of the queue to read from.
         :param int count: The maximum amount of messages to retrieve.
+        :param int visibility_timeout: The number of seconds the message should remain invisible to other queue readers.
         """
         queue = self.get_queue(queue_name)
         result = []
@@ -288,7 +289,7 @@ class SQSBackend(BaseBackend):
                 if not message.visible:
                     continue
                 message.mark_received(
-                    visibility_timeout=queue.visibility_timeout
+                    visibility_timeout=visibility_timeout
                 )
                 result.append(message)
                 if len(result) >= count:

moto/sqs/responses.py

@@ -38,6 +38,18 @@ class SQSResponse(BaseResponse):
             queue_name = self.path.split("/")[-1]
         return queue_name
 
+    def _get_validated_visibility_timeout(self):
+        """
+        :raises ValueError: If specified visibility timeout exceeds MAXIMUM_VISIBILTY_TIMEOUT
+        :raises TypeError: If visibility timeout was not specified
+        """
+        visibility_timeout = int(self.querystring.get("VisibilityTimeout")[0])
+
+        if visibility_timeout > MAXIMUM_VISIBILTY_TIMEOUT:
+            raise ValueError
+
+        return visibility_timeout
+
     def create_queue(self):
         queue_name = self.querystring.get("QueueName")[0]
         queue = self.sqs_backend.create_queue(queue_name, visibility_timeout=self.attribute.get('VisibilityTimeout'),
@@ -63,12 +75,11 @@ class SQSResponse(BaseResponse):
     def change_message_visibility(self):
         queue_name = self._get_queue_name()
         receipt_handle = self.querystring.get("ReceiptHandle")[0]
-        visibility_timeout = int(self.querystring.get("VisibilityTimeout")[0])
 
-        if visibility_timeout > MAXIMUM_VISIBILTY_TIMEOUT:
-            return "Invalid request, maximum visibility timeout is {0}".format(
-                MAXIMUM_VISIBILTY_TIMEOUT
-            ), dict(status=400)
+        try:
+            visibility_timeout = self._get_validated_visibility_timeout()
+        except ValueError:
+            return ERROR_MAX_VISIBILITY_TIMEOUT_RESPONSE, dict(status=400)
 
         try:
             self.sqs_backend.change_message_visibility(
@@ -233,7 +244,14 @@ class SQSResponse(BaseResponse):
         except TypeError:
             wait_time = queue.wait_time_seconds
 
-        messages = self.sqs_backend.receive_messages(queue_name, message_count, wait_time)
+        try:
+            visibility_timeout = self._get_validated_visibility_timeout()
+        except TypeError:
+            visibility_timeout = queue.visibility_timeout
+        except ValueError:
+            return ERROR_MAX_VISIBILITY_TIMEOUT_RESPONSE, dict(status=400)
+
+        messages = self.sqs_backend.receive_messages(queue_name, message_count, wait_time, visibility_timeout)
         template = self.response_template(RECEIVE_MESSAGE_RESPONSE)
         output = template.render(messages=messages)
         return output
@@ -436,3 +454,5 @@ ERROR_TOO_LONG_RESPONSE = """<ErrorResponse xmlns="http://queue.amazonaws.com/do
     </Error>
     <RequestId>6fde8d1e-52cd-4581-8cd9-c512f4c64223</RequestId>
 </ErrorResponse>"""
+
+ERROR_MAX_VISIBILITY_TIMEOUT_RESPONSE = "Invalid request, maximum visibility timeout is {0}".format(MAXIMUM_VISIBILTY_TIMEOUT)

tests/test_sqs/test_sqs.py

@@ -203,6 +203,23 @@ def test_message_becomes_inflight_when_received():
     queue.count().should.equal(1)
 
 
+@mock_sqs
+def test_receive_message_with_explicit_visibility_timeout():
+    conn = boto.connect_sqs('the_key', 'the_secret')
+    queue = conn.create_queue("test-queue", visibility_timeout=60)
+    queue.set_message_class(RawMessage)
+
+    body_one = 'this is another test message'
+    queue.write(queue.new_message(body_one))
+
+    queue.count().should.equal(1)
+    messages = conn.receive_message(queue, number_messages=1, visibility_timeout=0)
+
+    assert len(messages) == 1
+
+    # Message should remain visible
+    queue.count().should.equal(1)
+
 @mock_sqs
 def test_change_message_visibility():
     conn = boto.connect_sqs('the_key', 'the_secret')