refactor: incorporate will_accept() checks into publish() (#108)

Author: Gurov Ilya

google/cloud/pubsub_v1/publisher/_batch/base.py

@@ -109,32 +109,6 @@ def status(self):
  """
  raise NotImplementedError
 
- def will_accept(self, message):
- """Return True if the batch is able to accept the message.
-
- In concurrent implementations, the attributes on the current batch
- may be modified by other workers. With this in mind, the caller will
- likely want to hold a lock that will make sure the state remains
- the same after the "will accept?" question is answered.
-
- Args:
- message (~.pubsub_v1.types.PubsubMessage): The Pub/Sub message.
-
- Returns:
- bool: Whether this batch can accept the message.
- """
- # If this batch is not accepting messages generally, return False.
- if self.status != BatchStatus.ACCEPTING_MESSAGES:
- return False
-
- # If this message will make the batch exceed the ``max_messages``
- # setting, return False.
- if len(self.messages) >= self.settings.max_messages:
- return False
-
- # Okay, everything is good.
- return True
-
  def cancel(self, cancellation_reason):
  """Complete pending futures with an exception.
 

google/cloud/pubsub_v1/publisher/_batch/thread.py

@@ -333,8 +333,8 @@ def publish(self, message):
  self._status != base.BatchStatus.ERROR
  ), "Publish after stop() or publish error."
 
- if not self.will_accept(message):
- return future
+ if self.status != base.BatchStatus.ACCEPTING_MESSAGES:
+ return
 
  size_increase = types.PublishRequest(messages=[message]).ByteSize()
 

tests/unit/pubsub_v1/publisher/batch/test_base.py

@@ -46,33 +46,3 @@ def test_len():
  assert len(batch) == 0
  batch.publish(types.PubsubMessage(data=b"foo"))
  assert len(batch) == 1
-
-
-def test_will_accept():
- batch = create_batch(status=BatchStatus.ACCEPTING_MESSAGES)
- message = types.PubsubMessage()
- assert batch.will_accept(message) is True
-
-
-def test_will_accept_oversize():
- batch = create_batch(
- settings=types.BatchSettings(max_bytes=10),
- status=BatchStatus.ACCEPTING_MESSAGES,
- )
- message = types.PubsubMessage(data=b"abcdefghijklmnopqrstuvwxyz")
- assert batch.will_accept(message) is True
-
-
-def test_will_not_accept_status():
- batch = create_batch(status="talk to the hand")
- message = types.PubsubMessage()
- assert batch.will_accept(message) is False
-
-
-def test_will_not_accept_number():
- batch = create_batch(
- settings=types.BatchSettings(max_messages=-1),
- status=BatchStatus.ACCEPTING_MESSAGES,
- )
- message = types.PubsubMessage(data=b"abc")
- assert batch.will_accept(message) is False

tests/unit/pubsub_v1/publisher/batch/test_thread.py

@@ -287,18 +287,56 @@ def test_publish_updating_batch_size():
  assert batch.size > 0 # I do not always trust protobuf.
 
 
-def test_publish_not_will_accept():
+def test_publish():
+ batch = create_batch()
+ message = types.PubsubMessage()
+ future = batch.publish(message)
+
+ assert len(batch.messages) == 1
+ assert batch._futures == [future]
+
+
+def test_publish_max_messages_zero():
  batch = create_batch(topic="topic_foo", max_messages=0)
- base_request_size = types.PublishRequest(topic="topic_foo").ByteSize()
 
- # Publish the message.
  message = types.PubsubMessage(data=b"foobarbaz")
+ with mock.patch.object(batch, "commit") as commit:
+ future = batch.publish(message)
+
+ assert future is not None
+ assert len(batch.messages) == 1
+ assert batch._futures == [future]
+ commit.assert_called_once()
+
+
+def test_publish_max_messages_enforced():
+ batch = create_batch(topic="topic_foo", max_messages=1)
+
+ message = types.PubsubMessage(data=b"foobarbaz")
+ message2 = types.PubsubMessage(data=b"foobarbaz2")
+
+ future = batch.publish(message)
+ future2 = batch.publish(message2)
+
+ assert future is not None
+ assert future2 is None
+ assert len(batch.messages) == 1
+ assert len(batch._futures) == 1
+
+
+def test_publish_max_bytes_enforced():
+ batch = create_batch(topic="topic_foo", max_bytes=15)
+
+ message = types.PubsubMessage(data=b"foobarbaz")
+ message2 = types.PubsubMessage(data=b"foobarbaz2")
+
  future = batch.publish(message)
+ future2 = batch.publish(message2)
 
- assert future is None
- assert batch.size == base_request_size
- assert batch.messages == []
- assert batch._futures == []
+ assert future is not None
+ assert future2 is None
+ assert len(batch.messages) == 1
+ assert len(batch._futures) == 1
 
 
 def test_publish_exceed_max_messages():

tests/unit/pubsub_v1/publisher/test_publisher_client.py

@@ -135,7 +135,6 @@ def test_publish():
  batch = mock.Mock(spec=client._batch_class)
 
  # Set the mock up to claim indiscriminately that it accepts all messages.
- batch.will_accept.return_value = True
  batch.publish.side_effect = (future1, future2)
 
  topic = "topic/path"
@@ -169,7 +168,6 @@ def test_publish_error_exceeding_flow_control_limits():
  client = publisher.Client(credentials=creds, publisher_options=publisher_options)
 
  mock_batch = mock.Mock(spec=client._batch_class)
- mock_batch.will_accept.return_value = True
  topic = "topic/path"
  client._set_batch(topic, mock_batch)
 
@@ -216,7 +214,6 @@ def test_publish_attrs_bytestring():
  # Use a mock in lieu of the actual batch class.
  batch = mock.Mock(spec=client._batch_class)
  # Set the mock up to claim indiscriminately that it accepts all messages.
- batch.will_accept.return_value = True
 
  topic = "topic/path"
  client._set_batch(topic, batch)
@@ -431,7 +428,6 @@ def test_publish_with_ordering_key():
  future1.add_done_callback = mock.Mock(spec=["__call__"])
  future2.add_done_callback = mock.Mock(spec=["__call__"])
 
- batch.will_accept.return_value = True
  batch.publish.side_effect = (future1, future2)
 
  topic = "topic/path"