github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/gcp/tests/pubsub_matcher.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""PubSub verifier used for end-to-end test."""

# pytype: skip-file

import logging
import time
from collections import Counter

from hamcrest.core.base_matcher import BaseMatcher

from apache_beam.io.gcp.pubsub import PubsubMessage

__all__ = ['PubSubMessageMatcher']

# Protect against environments where pubsub library is not available.
try:
  from google.cloud import pubsub
except ImportError:
  pubsub = None

DEFAULT_TIMEOUT = 5 * 60
DEFAULT_SLEEP_TIME = 1
DEFAULT_MAX_MESSAGES_IN_ONE_PULL = 50
DEFAULT_PULL_TIMEOUT = 30.0

_LOGGER = logging.getLogger(__name__)


class PubSubMessageMatcher(BaseMatcher):
  """Matcher that verifies messages from given subscription.

  This matcher can block the test and keep pulling messages from given
  subscription until all expected messages are shown or timeout.
  """
  def __init__(
      self,
      project,
      sub_name,
      expected_msg=None,
      expected_msg_len=None,
      timeout=DEFAULT_TIMEOUT,
      with_attributes=False,
      strip_attributes=None,
      sleep_time=DEFAULT_SLEEP_TIME,
      max_messages_in_one_pull=DEFAULT_MAX_MESSAGES_IN_ONE_PULL,
      pull_timeout=DEFAULT_PULL_TIMEOUT):
    """Initialize PubSubMessageMatcher object.

    Args:
      project: A name string of project.
      sub_name: A name string of subscription which is attached to output.
      expected_msg: A string list that contains expected message data pulled
        from the subscription. See also: with_attributes.
      expected_msg_len: Number of expected messages pulled from the
        subscription.
      timeout: Timeout in seconds to wait for all expected messages appears.
      with_attributes: If True, will match against both message data and
        attributes. If True, expected_msg should be a list of ``PubsubMessage``
        objects. Otherwise, it should be a list of ``bytes``.
      strip_attributes: List of strings. If with_attributes==True, strip the
        attributes keyed by these values from incoming messages.
        If a key is missing, will add an attribute with an error message as
        value to prevent a successful match.
      sleep_time: Time in seconds between which the pulls from pubsub are done.
      max_messages_in_one_pull: Maximum number of messages pulled from pubsub
        at once.
      pull_timeout: Time in seconds after which the pull from pubsub is repeated
    """
    if pubsub is None:
      raise ImportError('PubSub dependencies are not installed.')
    if not project:
      raise ValueError('Invalid project %s.' % project)
    if not sub_name:
      raise ValueError('Invalid subscription %s.' % sub_name)
    if not expected_msg_len and not expected_msg:
      raise ValueError(
          'Required expected_msg: {} or expected_msg_len: {}.'.format(
              expected_msg, expected_msg_len))
    if expected_msg and not isinstance(expected_msg, list):
      raise ValueError('Invalid expected messages %s.' % expected_msg)
    if expected_msg_len and not isinstance(expected_msg_len, int):
      raise ValueError('Invalid expected messages %s.' % expected_msg_len)

    self.project = project
    self.sub_name = sub_name
    self.expected_msg = expected_msg
    self.expected_msg_len = expected_msg_len or len(self.expected_msg)
    self.timeout = timeout
    self.messages = None
    self.messages_all_details = None
    self.with_attributes = with_attributes
    self.strip_attributes = strip_attributes
    self.sleep_time = sleep_time
    self.max_messages_in_one_pull = max_messages_in_one_pull
    self.pull_timeout = pull_timeout

  def _matches(self, _):
    if self.messages is None:
      self.messages, self.messages_all_details = self._wait_for_messages(
          self.expected_msg_len, self.timeout)
    if self.expected_msg:
      return Counter(self.messages) == Counter(self.expected_msg)
    else:
      return len(self.messages) == self.expected_msg_len

  def _wait_for_messages(self, expected_num, timeout):
    """Wait for messages from given subscription."""
    total_messages = []
    total_messages_all_details = []

    sub_client = pubsub.SubscriberClient()
    start_time = time.time()
    while time.time() - start_time <= timeout:
      response = sub_client.pull(
          subscription=self.sub_name,
          max_messages=self.max_messages_in_one_pull,
          timeout=self.pull_timeout)
      for rm in response.received_messages:
        msg = PubsubMessage._from_message(rm.message)
        full_message = (
            msg.data,
            msg.attributes,
            msg.attributes,
            msg.publish_time,
            msg.ordering_key)
        if not self.with_attributes:
          total_messages.append(msg.data)
          total_messages_all_details.append(full_message)
          continue

        if self.strip_attributes:
          for attr in self.strip_attributes:
            try:
              del msg.attributes[attr]
            except KeyError:
              msg.attributes[attr] = (
                  'PubSubMessageMatcher error: '
                  'expected attribute not found.')
        total_messages.append(msg)
        total_messages_all_details.append(full_message)

      ack_ids = [rm.ack_id for rm in response.received_messages]
      if ack_ids:
        sub_client.acknowledge(subscription=self.sub_name, ack_ids=ack_ids)
      if len(total_messages) >= expected_num:
        break
      time.sleep(self.sleep_time)

    if time.time() - start_time > timeout:
      _LOGGER.error(
          'Timeout after %d sec. Received %d messages from %s.',
          timeout,
          len(total_messages),
          self.sub_name)
    return total_messages, total_messages_all_details

  def describe_to(self, description):
    description.append_text('Expected %d messages.' % self.expected_msg_len)

  def describe_mismatch(self, _, mismatch_description):
    c_expected = Counter(self.expected_msg)
    c_actual = Counter(self.messages)
    mismatch_description.append_text("Got %d messages. " % (len(self.messages)))
    if self.expected_msg:
      expected = (c_expected - c_actual).items()
      unexpected = (c_actual - c_expected).items()
      unexpected_keys = [repr(item[0]) for item in unexpected]
      if self.with_attributes:
        unexpected_all_details = [
            x for x in self.messages_all_details
            if 'PubsubMessage(%s, %s)' % (repr(x[0]), x[1]) in unexpected_keys
        ]
      else:
        unexpected_all_details = [
            x for x in self.messages_all_details
            if repr(x[0]) in unexpected_keys
        ]
      mismatch_description.append_text(
          "Diffs (item, count):\n"
          "  Expected but not in actual: %s\n"
          "  Unexpected: %s\n"
          "  Unexpected (with all details): %s" %
          (expected, unexpected, unexpected_all_details))
    if self.with_attributes and self.strip_attributes:
      mismatch_description.append_text(
          '\n  Stripped attributes: %r' % self.strip_attributes)