def test__send_offset_requests_multiple_nodes(fetcher, mocker):
tp1 = TopicPartition("topic_send_offset", 1)
tp2 = TopicPartition("topic_send_offset", 2)
tp3 = TopicPartition("topic_send_offset", 3)
tp4 = TopicPartition("topic_send_offset", 4)
mocked_send = mocker.patch.object(fetcher, "_send_offset_request")
send_futures = []
def send_side_effect(node_id, timestamps):
f = Future()
send_futures.append((node_id, timestamps, f))
return f
mocked_send.side_effect = send_side_effect
mocked_leader = mocker.patch.object(fetcher._client.cluster,
"leader_for_partition")
mocked_leader.side_effect = itertools.cycle([0, 1])
# -- All node succeeded case
tss = OrderedDict([(tp1, 0), (tp2, 0), (tp3, 0), (tp4, 0)])
fut = fetcher._send_offset_requests(tss)
assert not fut.is_done
assert mocked_send.call_count == 2
req_by_node = {}
second_future = None
for node, timestamps, f in send_futures:
req_by_node[node] = timestamps
if node == 0:
# Say tp3 does not have any messages so it's missing
f.success({tp1: (11, 1001)})
else:
second_future = f
assert req_by_node == {0: {tp1: 0, tp3: 0}, 1: {tp2: 0, tp4: 0}}
# We only resolved 1 future so far, so result future is not yet ready
assert not fut.is_done
second_future.success({tp2: (12, 1002), tp4: (14, 1004)})
assert fut.succeeded()
assert fut.value == {tp1: (11, 1001), tp2: (12, 1002), tp4: (14, 1004)}
# -- First succeeded second not
del send_futures[:]
fut = fetcher._send_offset_requests(tss)
assert len(send_futures) == 2
send_futures[0][2].success({tp1: (11, 1001)})
send_futures[1][2].failure(UnknownTopicOrPartitionError(tp1))
assert fut.failed()
assert isinstance(fut.exception, UnknownTopicOrPartitionError)
# -- First fails second succeeded
del send_futures[:]
fut = fetcher._send_offset_requests(tss)
assert len(send_futures) == 2
send_futures[0][2].failure(UnknownTopicOrPartitionError(tp1))
send_futures[1][2].success({tp1: (11, 1001)})
assert fut.failed()
assert isinstance(fut.exception, UnknownTopicOrPartitionError)
def _get_leader_for_partition(self, topic, partition):
"""
Returns the leader for a partition or None if the partition exists
but has no leader.
UnknownTopicOrPartitionError will be raised if the topic or partition
is not part of the metadata.
LeaderNotAvailableError is raised if server has metadata, but there is
no current leader
"""
key = TopicPartition(topic, partition)
# Use cached metadata if it is there
if self.topics_to_brokers.get(key) is not None:
return self.topics_to_brokers[key]
# Otherwise refresh metadata
# If topic does not already exist, this will raise
# UnknownTopicOrPartitionError if not auto-creating
# LeaderNotAvailableError otherwise until partitions are created
self.load_metadata_for_topics(topic)
# If the partition doesn't actually exist, raise
if partition not in self.topic_partitions.get(topic, []):
raise UnknownTopicOrPartitionError(key)
# If there's no leader for the partition, raise
leader = self.topic_partitions[topic][partition]
if leader == -1:
raise LeaderNotAvailableError((topic, partition))
# Otherwise return the BrokerMetadata
return self.brokers[leader]
def send(
self,
topic=None,
value=None,
key=None,
headers=None,
partition=None,
timestamp_ms=None,
bootstrap_servers=None,
):
producer = self.get_producer(bootstrap_servers if bootstrap_servers
is not None else [DEFAULT_FLAG])
if producer is None:
raise Exception("no available producer")
topic = topic if topic is not None else self.topic
if topic is None:
raise Exception("no topic")
if (len(producer._metadata._partitions) > 0
and topic not in producer._metadata._partitions):
if topic in self.not_exist_topics:
try:
producer._wait_on_metadata(topic, 1)
except Exception:
return Future().failure(
UnknownTopicOrPartitionError(topic))
else:
self.not_exist_topics[topic] = 1
return producer.send(
topic,
value=value,
key=key,
headers=headers,
partition=partition,
timestamp_ms=timestamp_ms,
)