<commit_before>"""
Message consumer.

"""
from boto3 import client
from microcosm.api import defaults


class SQSConsumer(object):
    """
    Consume message from a (single) SQS queue.

    """
    def __init__(
        self,
        sqs_client,
        sqs_envelope,
        sqs_queue_url,
        limit,
        wait_seconds,
        visibility_timeout_seconds,
    ):
        self.sqs_client = sqs_client
        self.sqs_envelope = sqs_envelope
        self.sqs_queue_url = sqs_queue_url
        self.limit = limit
        self.wait_seconds = wait_seconds
        self.visibility_timeout_seconds = visibility_timeout_seconds

    def consume(self):
        """
        Consume a batch of messages.

        :returns: a list of `SQSMessage`
        """
        return [
            self.sqs_envelope.parse_raw_message(self, raw_message)
            for raw_message in self.sqs_client.receive_message(
                QueueUrl=self.sqs_queue_url,
                MaxNumberOfMessages=self.limit,
                WaitTimeSeconds=self.wait_seconds,
            ).get("Messages", [])
        ]

    def ack(self, message):
        """
        Acknowledge that a message was processed successfully.

        Deletes the message from the queue.

        """
        self.sqs_client.delete_message(
            QueueUrl=self.sqs_queue_url,
            ReceiptHandle=message.receipt_handle,
        )

    def nack(self, message):
        """
        Acknowledge that a message was NOT processed successfully.

        Sets the visibility timeout if present

        """
        if self.visibility_timeout_seconds is not None:
            self.sqs_client.change_message_visibility(
                QueueUrl=self.sqs_queue_url,
                ReceiptHandle=message.receipt_handle,
                VisibilityTimeout=self.visibility_timeout_seconds,
            )


@defaults(
    # SQS will not return more than ten messages at a time
    limit=10,
    # SQS will only return a few messages at time unless long polling is enabled (>0)
    wait_seconds=1,
)
def configure_sqs_consumer(graph):
    """
    Configure an SQS consumer.

    """
    sqs_queue_url = graph.config.sqs_consumer.sqs_queue_url
    limit = graph.config.sqs_consumer.limit
    wait_seconds = graph.config.sqs_consumer.wait_seconds
    visibility_timeout_seconds = graph.config.sqs_consumer.visibility_timeout_seconds

    if graph.metadata.testing:
        from mock import MagicMock
        sqs_client = MagicMock()
    else:
        sqs_client = client("sqs")

    return SQSConsumer(
        sqs_client=sqs_client,
        sqs_envelope=graph.sqs_envelope,
        sqs_queue_url=sqs_queue_url,
        limit=limit,<commit_msg>cast visibility timeout as int of not None<commit_after>"""
Message consumer.

"""
from boto3 import client
from microcosm.api import defaults


class SQSConsumer(object):
    """
    Consume message from a (single) SQS queue.

    """
    def __init__(
        self,
        sqs_client,
        sqs_envelope,
        sqs_queue_url,
        limit,
        wait_seconds,
        visibility_timeout_seconds,
    ):
        self.sqs_client = sqs_client
        self.sqs_envelope = sqs_envelope
        self.sqs_queue_url = sqs_queue_url
        self.limit = limit
        self.wait_seconds = wait_seconds
        self.visibility_timeout_seconds = visibility_timeout_seconds

    def consume(self):
        """
        Consume a batch of messages.

        :returns: a list of `SQSMessage`
        """
        return [
            self.sqs_envelope.parse_raw_message(self, raw_message)
            for raw_message in self.sqs_client.receive_message(
                QueueUrl=self.sqs_queue_url,
                MaxNumberOfMessages=self.limit,
                WaitTimeSeconds=self.wait_seconds,
            ).get("Messages", [])
        ]

    def ack(self, message):
        """
        Acknowledge that a message was processed successfully.

        Deletes the message from the queue.

        """
        self.sqs_client.delete_message(
            QueueUrl=self.sqs_queue_url,
            ReceiptHandle=message.receipt_handle,
        )

    def nack(self, message):
        """
        Acknowledge that a message was NOT processed successfully.

        Sets the visibility timeout if present

        """
        if self.visibility_timeout_seconds is not None:
            self.sqs_client.change_message_visibility(
                QueueUrl=self.sqs_queue_url,
                ReceiptHandle=message.receipt_handle,
                VisibilityTimeout=self.visibility_timeout_seconds,
            )


@defaults(
    # SQS will not return more than ten messages at a time
    limit=10,
    # SQS will only return a few messages at time unless long polling is enabled (>0)
    wait_seconds=1,
)
def configure_sqs_consumer(graph):
    """
    Configure an SQS consumer.

    """
    sqs_queue_url = graph.config.sqs_consumer.sqs_queue_url
    limit = graph.config.sqs_consumer.limit
    wait_seconds = graph.config.sqs_consumer.wait_seconds
    visibility_timeout_seconds = graph.config.sqs_consumer.visibility_timeout_seconds

    if visibility_timeout_seconds:
        visibility_timeout_seconds = int(visibility_timeout_seconds)

    if graph.metadata.testing:
        from mock import MagicMock
        sqs_client = MagicMock()
    else:
        sqs_client = client("sqs")

    return SQSConsumer(
        sqs_client=sqs_client,
        sqs_envelope=graph.sqs_envelope,
        sqs_queue_url=sqs_queue_url,
        limit=limit,
{'fraction_non_alphanumeric': 0.0713768115942029, 'fraction_numerical': 0.0018115942028985507, 'mean_word_length': 2.2824019024970275, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7561309', 'n_tokens_mistral': 1679, 'n_tokens_neox': 1584, 'n_words': 348}