Moving will_accept() check out of PublisherClient.batch() factory. (googleapis#4613)

dhermes · web-flow · commit 72cbe413a133 · 2017-12-18T13:03:55.000-08:00
This is the first stage of moving that check directly into
`Batch.publish()`. This check **must** occur when the message
is being published because concurrent access to the messages
and futures on a `Batch` make it impossible to reliably do
an LBYL "will accept?" check.
diff --git a/pubsub/google/cloud/pubsub_v1/publisher/client.py b/pubsub/google/cloud/pubsub_v1/publisher/client.py
@@ -98,30 +98,34 @@ def target(self):
         """
         return publisher_client.PublisherClient.SERVICE_ADDRESS
 
-    def batch(self, topic, message, create=True, autocommit=True):
+    def batch(self, topic, create=False, autocommit=True):
         """Return the current batch for the provided topic.
 
-        This will create a new batch only if no batch currently exists.
+        This will create a new batch if ``create=True`` or if no batch
+        currently exists.
 
         Args:
             topic (str): A string representing the topic.
-            message (~google.cloud.pubsub_v1.types.PubsubMessage): The message
-                that will be committed.
-            create (bool): Whether to create a new batch if no batch is
-                found. Defaults to True.
-            autocommit (bool): Whether to autocommit this batch.
-                This is primarily useful for debugging.
+            create (bool): Whether to create a new batch. Defaults to
+                :data:`False`. If :data:`True`, this will create a new batch
+                even if one already exists.
+            autocommit (bool): Whether to autocommit this batch. This is
+                primarily useful for debugging and testing, since it allows
+                the caller to avoid some side effects that batch creation
+                might have (e.g. spawning a worker to publish a batch).
 
         Returns:
             ~.pubsub_v1.batch.Batch: The batch object.
         """
         # If there is no matching batch yet, then potentially create one
         # and place it on the batches dictionary.
         with self._batch_lock:
-            batch = self._batches.get(topic, None)
-            if batch is None or not batch.will_accept(message):
-                if not create:
-                    return None
+            if not create:
+                batch = self._batches.get(topic)
+                if batch is None:
+                    create = True
+
+            if create:
                 batch = self._batch_class(
                     autocommit=autocommit,
                     client=self,
@@ -130,7 +134,6 @@ def batch(self, topic, message, create=True, autocommit=True):
                 )
                 self._batches[topic] = batch
 
-        # Simply return the appropriate batch.
         return batch
 
     def publish(self, topic, data, **attrs):
@@ -190,4 +193,12 @@ def publish(self, topic, data, **attrs):
         message = types.PubsubMessage(data=data, attributes=attrs)
 
         # Delegate the publishing to the batch.
-        return self.batch(topic, message=message).publish(message)
+        batch = self.batch(topic)
+        future = None
+        while future is None:
+            if batch.will_accept(message):
+                future = batch.publish(message)
+            else:
+                batch = self.batch(topic, create=True)
+
+        return future
diff --git a/pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py b/pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py
@@ -51,100 +51,148 @@ def test_init_emulator(monkeypatch):
     assert channel.target().decode('utf8') == '/foo/bar/'
 
 
-def test_batch_accepting():
-    """Establish that an existing batch is returned if it accepts messages."""
+def test_batch_create():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)
-    message = types.PubsubMessage(data=b'foo')
 
-    # At first, there are no batches, so this should return a new batch
-    # which is also saved to the object.
-    ante = len(client._batches)
-    batch = client.batch('topic_name', message, autocommit=False)
-    assert len(client._batches) == ante + 1
-    assert batch is client._batches['topic_name']
+    assert len(client._batches) == 0
+    topic = 'topic/path'
+    batch = client.batch(topic, autocommit=False)
+    assert client._batches == {topic: batch}
+
+
+def test_batch_exists():
+    creds = mock.Mock(spec=credentials.Credentials)
+    client = publisher.Client(credentials=creds)
+
+    topic = 'topic/path'
+    client._batches[topic] = mock.sentinel.batch
 
     # A subsequent request should return the same batch.
-    batch2 = client.batch('topic_name', message, autocommit=False)
-    assert batch is batch2
-    assert batch2 is client._batches['topic_name']
+    batch = client.batch(topic, autocommit=False)
+    assert batch is mock.sentinel.batch
+    assert client._batches == {topic: batch}
 
 
-def test_batch_without_autocreate():
+def test_batch_create_and_exists():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)
-    message = types.PubsubMessage(data=b'foo')
 
-    # If `create=False` is sent, then when the batch is not found, None
-    # is returned instead.
-    ante = len(client._batches)
-    batch = client.batch('topic_name', message, create=False)
-    assert batch is None
-    assert len(client._batches) == ante
+    topic = 'topic/path'
+    client._batches[topic] = mock.sentinel.batch
+
+    # A subsequent request should return the same batch.
+    batch = client.batch(topic, create=True, autocommit=False)
+    assert batch is not mock.sentinel.batch
+    assert client._batches == {topic: batch}
 
 
 def test_publish():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)
 
-    # Use a mock in lieu of the actual batch class; set the mock up to claim
-    # indiscriminately that it accepts all messages.
+    # 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
-    client._batches['topic_name'] = batch
-
-    # Begin publishing.
-    client.publish('topic_name', b'spam')
-    client.publish('topic_name', b'foo', bar='baz')
 
-    # The batch's publish method should have been called twice.
-    assert batch.publish.call_count == 2
+    topic = 'topic/path'
+    client._batches[topic] = batch
 
-    # In both cases
-    # The first call should correspond to the first message.
-    _, args, _ = batch.publish.mock_calls[0]
-    assert args[0].data == b'spam'
-    assert not args[0].attributes
+    # Begin publishing.
+    client.publish(topic, b'spam')
+    client.publish(topic, b'foo', bar='baz')
 
-    # The second call should correspond to the second message.
-    _, args, _ = batch.publish.mock_calls[1]
-    assert args[0].data == b'foo'
-    assert args[0].attributes == {u'bar': u'baz'}
+    # Check mock.
+    batch.publish.assert_has_calls(
+        [
+            mock.call(types.PubsubMessage(data=b'spam')),
+            mock.call(types.PubsubMessage(
+                data=b'foo',
+                attributes={'bar': 'baz'},
+            )),
+        ],
+    )
 
 
 def test_publish_data_not_bytestring_error():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)
+    topic = 'topic/path'
     with pytest.raises(TypeError):
-        client.publish('topic_name', u'This is a text string.')
+        client.publish(topic, u'This is a text string.')
     with pytest.raises(TypeError):
-        client.publish('topic_name', 42)
+        client.publish(topic, 42)
 
 
 def test_publish_attrs_bytestring():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)
 
-    # Use a mock in lieu of the actual batch class; set the mock up to claim
-    # indiscriminately that it accepts all messages.
+    # 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
-    client._batches['topic_name'] = batch
+
+    topic = 'topic/path'
+    client._batches[topic] = batch
+
+    # Begin publishing.
+    client.publish(topic, b'foo', bar=b'baz')
+
+    # The attributes should have been sent as text.
+    batch.publish.assert_called_once_with(
+        types.PubsubMessage(
+            data=b'foo',
+            attributes={'bar': u'baz'},
+        ),
+    )
+
+
+def test_publish_new_batch_needed():
+    creds = mock.Mock(spec=credentials.Credentials)
+    client = publisher.Client(credentials=creds)
+
+    # Use mocks in lieu of the actual batch class.
+    batch1 = mock.Mock(spec=client._batch_class)
+    batch2 = mock.Mock(spec=client._batch_class)
+    # Set the first mock up to claim indiscriminately that it rejects all
+    # messages and the second accepts all.
+    batch1.will_accept.return_value = False
+    batch2.will_accept.return_value = True
+
+    topic = 'topic/path'
+    client._batches[topic] = batch1
+
+    # Actually mock the batch class now.
+    batch_class = mock.Mock(spec=(), return_value=batch2)
+    client._batch_class = batch_class
 
     # Begin publishing.
-    client.publish('topic_name', b'foo', bar=b'baz')
+    client.publish(topic, b'foo', bar=b'baz')
+    batch_class.assert_called_once_with(
+        autocommit=True,
+        client=client,
+        settings=client.batch_settings,
+        topic=topic,
+    )
 
     # The attributes should have been sent as text.
-    _, args, _ = batch.publish.mock_calls[0]
-    assert args[0].data == b'foo'
-    assert args[0].attributes == {u'bar': u'baz'}
+    batch1.publish.assert_not_called()
+    batch2.publish.assert_called_once_with(
+        types.PubsubMessage(
+            data=b'foo',
+            attributes={'bar': u'baz'},
+        ),
+    )
 
 
 def test_publish_attrs_type_error():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)
+    topic = 'topic/path'
     with pytest.raises(TypeError):
-        client.publish('topic_name', b'foo', answer=42)
+        client.publish(topic, b'foo', answer=42)
 
 
 def test_gapic_instance_method():
